diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index 6f1bd344540..b34585f2532 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -17,12 +17,14 @@ steps: - tests/kernels/test_awq_int4_to_int8.py - tests/kernels/quantization/test_cpu_fp8_scaled_mm.py - tests/kernels/mamba/cpu/test_cpu_gdn_ops.py + - tests/kernels/mamba/test_cpu_short_conv.py commands: - | bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m " pytest -x -v -s tests/kernels/attention/test_cpu_attn.py pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py pytest -x -v -s tests/kernels/moe/test_cpu_quant_fused_moe.py + pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py pytest -x -v -s tests/kernels/test_onednn.py pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py diff --git a/.buildkite/intel_jobs/misc_intel.yaml b/.buildkite/intel_jobs/misc_intel.yaml index ec3d34c20fe..c9bebb7bc14 100644 --- a/.buildkite/intel_jobs/misc_intel.yaml +++ b/.buildkite/intel_jobs/misc_intel.yaml @@ -76,6 +76,30 @@ steps: pytest -v -s v1/sample/test_logprobs.py && pytest -v -s v1/sample/test_logprobs_e2e.py' +- label: Basic Models Tests (Initialization) + timeout_in_minutes: 60 + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 24+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/ + - tests/models/test_initialization.py + - tests/models/registry.py + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'export VLLM_XPU_FUSED_MOE_USE_REF=1 && + cd tests && + pytest -v -s models/test_initialization.py::test_can_initialize_large_subset[Eagle3MiniMaxM2ForCausalLM]' + - label: XPU CPU Offload timeout_in_minutes: 60 device: intel_gpu diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh index 9db805c4585..3969c7f974e 100755 --- a/.buildkite/scripts/hardware_ci/run-amd-test.sh +++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh @@ -551,6 +551,7 @@ else fi docker run \ + -t -i \ --device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \ $RDMA_FLAGS \ --network=host \ diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh index 252eeeef8ce..2d11dd477ea 100755 --- a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh @@ -38,7 +38,9 @@ function cpu_tests() { pytest -x -v -s tests/kernels/attention/test_cpu_attn.py pytest -x -v -s tests/kernels/core/test_cpu_activation.py pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py - pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py" + pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py + pytest -x -v -s tests/kernels/moe/test_cpu_int4_moe.py + pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py" # skip tests requiring model downloads if HF_TOKEN is not set # due to rate-limits @@ -62,7 +64,6 @@ function cpu_tests() { set -e pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs" - # basic online serving docker exec cpu-test bash -c ' set -e diff --git a/.buildkite/scripts/run-multi-node-test.sh b/.buildkite/scripts/run-multi-node-test.sh index c0911f17b66..cf54986e015 100755 --- a/.buildkite/scripts/run-multi-node-test.sh +++ b/.buildkite/scripts/run-multi-node-test.sh @@ -109,7 +109,9 @@ run_nodes() { if [ "$node" -ne 0 ]; then docker exec -d "node$node" /bin/bash -c "cd $WORKING_DIR ; ${COMMANDS[$node]}" else - docker exec "node$node" /bin/bash -c "cd $WORKING_DIR ; ${COMMANDS[$node]}" + # Allocate a TTY (-t -i) for the foreground head node so its output + # keeps ANSI color in the Buildkite log (see run-amd-test.sh). + docker exec -t -i "node$node" /bin/bash -c "cd $WORKING_DIR ; ${COMMANDS[$node]}" fi done } diff --git a/.buildkite/scripts/run-rust-frontend-cargo-ci.sh b/.buildkite/scripts/run-rust-frontend-cargo-ci.sh index 42ab1fb543b..b19af8a2e5d 100755 --- a/.buildkite/scripts/run-rust-frontend-cargo-ci.sh +++ b/.buildkite/scripts/run-rust-frontend-cargo-ci.sh @@ -8,7 +8,12 @@ if [[ "$MODE" != "style-clippy" && "$MODE" != "test" ]]; then exit 2 fi -ROOT_DIR="$(git rev-parse --show-toplevel)" +if ROOT_DIR="$(git rev-parse --show-toplevel 2>/dev/null)"; then + : +else + SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" + ROOT_DIR="$(cd -- "${SCRIPT_DIR}/../.." && pwd -P)" +fi cd "$ROOT_DIR" export CARGO_TERM_COLOR="${CARGO_TERM_COLOR:-always}" diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 57c976030de..9c1b2077a7e 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -404,6 +404,36 @@ steps: - pytest -v -s transformers_utils - pytest -v -s config +#------------------------------------------------------------ mi250 · rust -----------------------------------------------------------# + +- label: Rust Frontend Cargo Style + Clippy # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + no_gpu: true + working_dir: "/vllm-workspace" + source_file_dependencies: + - rust/ + - rust-toolchain.toml + - .buildkite/test_areas/rust_frontend_cargo.yaml + - .buildkite/scripts/run-rust-frontend-cargo-ci.sh + commands: + - bash .buildkite/scripts/run-rust-frontend-cargo-ci.sh style-clippy + +- label: Rust Frontend Cargo Tests # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + no_gpu: true + working_dir: "/vllm-workspace" + source_file_dependencies: + - rust/ + - rust-toolchain.toml + - .buildkite/test_areas/rust_frontend_cargo.yaml + - .buildkite/scripts/run-rust-frontend-cargo-ci.sh + commands: + - bash .buildkite/scripts/run-rust-frontend-cargo-ci.sh test + #----------------------------------------------------------- mi250 · docker ----------------------------------------------------------# - label: Docker Build Metadata (ROCm) # TBD @@ -472,7 +502,7 @@ steps: commands: - TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)' - - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)' + - pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)' - pytest models/language -v -s -m 'distributed(num_gpus=2)' - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py - pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)' @@ -1301,7 +1331,7 @@ steps: #---------------------------------------------------------- mi300 · kernels ----------------------------------------------------------# - label: vLLM IR Tests # TBD - timeout_in_minutes: 30 + timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] agent_pool: mi300_1 optional: true @@ -1350,7 +1380,7 @@ steps: - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py - label: Kernels KDA Test # TBD - timeout_in_minutes: 30 + timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] agent_pool: mi300_1 optional: true @@ -1558,7 +1588,7 @@ steps: - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py - label: Model Runner V2 Pipeline Parallelism (4 GPUs) # TBD - timeout_in_minutes: 60 + timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] agent_pool: mi300_4 num_gpus: 4 @@ -1577,7 +1607,7 @@ steps: - pytest -v -s distributed/test_pp_cudagraph.py -k "not ray" - label: Model Runner V2 Spec Decode # TBD - timeout_in_minutes: 45 + timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] agent_pool: mi300_1 optional: true @@ -1604,7 +1634,7 @@ steps: timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] agent_pool: mi300_1 - parallelism: 2 + parallelism: 6 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/model_executor/models/ @@ -1637,10 +1667,10 @@ steps: source_file_dependencies: - vllm/ - tests/models/test_terratorch.py - - tests/models/test_transformers.py + - tests/models/transformers/test_backend.py - tests/models/test_registry.py commands: - - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py + - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py #----------------------------------------------------- mi300 · models / language -----------------------------------------------------# @@ -1860,7 +1890,7 @@ steps: - examples/ commands: - pip install --upgrade git+https://github.com/huggingface/transformers - - pytest -v -s tests/models/test_transformers.py + - pytest -v -s tests/models/transformers/test_backend.py - pytest -v -s tests/models/multimodal/test_mapping.py - python3 examples/basic/offline_inference/chat.py - python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl @@ -1915,7 +1945,7 @@ steps: - pytest -v -s plugins_tests/lora_resolvers # unit tests for in-tree lora resolver plugins - label: GGUF Plugin # TBD - timeout_in_minutes: 30 + timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] agent_pool: mi300_1 optional: true @@ -1930,6 +1960,118 @@ steps: - pip install "vllm-gguf-plugin >= 0.0.2" - pytest -v -s plugins_tests/gguf +#------------------------------------------------------- mi300 · rust_frontend -------------------------------------------------------# + +- label: Rust Frontend OpenAI Coverage # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - rust/ + - vllm/benchmarks/ + - vllm/entrypoints/openai/ + - vllm/entrypoints/serve/ + - vllm/v1/sample/ + - tests/utils.py + - tests/benchmarks/test_serve_cli.py + - tests/entrypoints/openai/chat_completion/test_chat_completion.py + - tests/entrypoints/openai/completion/test_shutdown.py + - tests/v1/sample/test_logprobs_e2e.py + - vllm/platforms/rocm.py + commands: + - export VLLM_USE_RUST_FRONTEND=1 + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)" + - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex" + - pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly" + - pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server" + +- label: Rust Frontend Serve Admin Coverage # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - rust/ + - vllm/entrypoints/openai/ + - vllm/entrypoints/serve/ + - vllm/v1/engine/ + - tests/utils.py + - tests/entrypoints/serve/disagg/test_serving_tokens.py + - tests/entrypoints/serve/instrumentator/test_basic.py + - tests/entrypoints/serve/instrumentator/test_metrics.py + - vllm/platforms/rocm.py + commands: + - export VLLM_USE_RUST_FRONTEND=1 + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load" + - pytest -v -s entrypoints/serve/disagg/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow" + - pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist" + +- label: Rust Frontend Core Correctness # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - rust/ + - vllm/entrypoints/openai/ + - tests/utils.py + - tests/entrypoints/openai/correctness/test_lmeval.py + - vllm/platforms/rocm.py + commands: + - export VLLM_USE_RUST_FRONTEND=1 + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine + +- label: Rust Frontend Tool Use # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - rust/ + - vllm/entrypoints/openai/ + - vllm/tool_parsers/ + - tests/utils.py + - tests/tool_use/ + - vllm/platforms/rocm.py + commands: + - export VLLM_USE_RUST_FRONTEND=1 + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice" + +- label: Rust Frontend Distributed # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - rust/ + - vllm/distributed/ + - vllm/engine/ + - vllm/executor/ + - vllm/v1/engine/ + - vllm/v1/worker/ + - tests/utils.py + - tests/v1/distributed/test_external_lb_dp.py + - tests/v1/distributed/test_hybrid_lb_dp.py + - tests/v1/distributed/test_internal_lb_dp.py + - vllm/platforms/rocm.py + commands: + - export VLLM_USE_RUST_FRONTEND=1 + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py -k "not 4 and not server_info" + - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py -k "not 4 and not server_info" + - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py -k "not 4 and not server_info" + #------------------------------------------------------- mi300 · quantization --------------------------------------------------------# - label: Quantization # TBD @@ -2439,6 +2581,59 @@ steps: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +- label: Hybrid SSM NixlConnector PD prefix cache test (2 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 + num_gpus: 2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - vllm/v1/core/sched/ + - vllm/v1/core/kv_cache_coordinator.py + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh + +- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 + num_gpus: 2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh + +- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 + num_gpus: 2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh + - label: V1 e2e (4 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -2729,15 +2924,19 @@ steps: - vllm/envs.py - examples/offline_inference/data_parallel.py - tests/distributed/test_context_parallel.py + - tests/distributed/test_rocm_aiter_custom_ar.py - tests/distributed/test_rocm_quick_reduce.py - tests/distributed/test_quick_all_reduce.py + - tests/v1/e2e/general/test_rocm_aiter_custom_ar.py - tests/v1/distributed/test_dbo.py - tests/utils.py commands: - pytest -v -s tests/distributed/test_context_parallel.py - - pytest -v -s tests/v1/distributed/test_dbo.py + - pytest -v -s tests/distributed/test_rocm_aiter_custom_ar.py + - pytest -v -s tests/v1/e2e/general/test_rocm_aiter_custom_ar.py - pytest -v -s tests/distributed/test_rocm_quick_reduce.py - pytest -v -s tests/distributed/test_quick_all_reduce.py + - pytest -v -s tests/v1/distributed/test_dbo.py #-------------------------------------------------------- mi355 · entrypoints --------------------------------------------------------# diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 8f47e7b5c21..4cf7774ebd4 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -150,9 +150,9 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt -- label: LM Eval Humming (A100 - TEMPORARY) - key: lm-eval-humming-a100 - timeout_in_minutes: 30 +- label: LM Eval Humming f16 (A100 - TEMPORARY) + key: lm-eval-humming-f16-a100 + timeout_in_minutes: 120 device: a100 optional: true num_devices: 1 @@ -160,13 +160,29 @@ steps: - vllm/model_executor/layers/quantization/humming.py - vllm/model_executor/layers/quantization/utils/humming_utils.py - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py - - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py + - vllm/model_executor/layers/fused_moe/oracle/ + - vllm/model_executor/kernels/linear/ commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt -- label: LM Eval Humming (H100 - TEMPORARY) - key: lm-eval-humming-h100 - timeout_in_minutes: 30 +- label: LM Eval Humming Act int8 (A100 - TEMPORARY) + key: lm-eval-humming-act-a100 + timeout_in_minutes: 120 + device: a100 + optional: true + num_devices: 1 + source_file_dependencies: + - vllm/model_executor/layers/quantization/humming.py + - vllm/model_executor/layers/quantization/utils/humming_utils.py + - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py + - vllm/model_executor/layers/fused_moe/oracle/ + - vllm/model_executor/kernels/linear/ + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt + +- label: LM Eval Humming f16 (H100 - TEMPORARY) + key: lm-eval-humming-f16-h100 + timeout_in_minutes: 120 device: h100 optional: true num_devices: 1 @@ -174,14 +190,30 @@ steps: - vllm/model_executor/layers/quantization/humming.py - vllm/model_executor/layers/quantization/utils/humming_utils.py - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py - - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py + - vllm/model_executor/layers/fused_moe/oracle/ + - vllm/model_executor/kernels/linear/ commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt -- label: LM Eval Humming (B200 - TEMPORARY) - key: lm-eval-humming-b200 - timeout_in_minutes: 30 +- label: LM Eval Humming Act fp8/int8 (H100 - TEMPORARY) + key: lm-eval-humming-act-h100 + timeout_in_minutes: 120 + device: h100 + optional: true + num_devices: 1 + source_file_dependencies: + - vllm/model_executor/layers/quantization/humming.py + - vllm/model_executor/layers/quantization/utils/humming_utils.py + - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py + - vllm/model_executor/layers/fused_moe/oracle/ + - vllm/model_executor/kernels/linear/ + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt + +- label: LM Eval Humming f16 (B200 - TEMPORARY) + key: lm-eval-humming-f16-b200 + timeout_in_minutes: 120 device: b200-k8s optional: true num_devices: 1 @@ -189,10 +221,26 @@ steps: - vllm/model_executor/layers/quantization/humming.py - vllm/model_executor/layers/quantization/utils/humming_utils.py - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py - - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py + - vllm/model_executor/layers/fused_moe/oracle/ + - vllm/model_executor/kernels/linear/ commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt + +- label: LM Eval Humming Act fp8/int8 (B200 - TEMPORARY) + key: lm-eval-humming-act-b200 + timeout_in_minutes: 120 + device: b200-k8s + optional: true + num_devices: 1 + source_file_dependencies: + - vllm/model_executor/layers/quantization/humming.py + - vllm/model_executor/layers/quantization/utils/humming_utils.py + - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py + - vllm/model_executor/layers/fused_moe/oracle/ + - vllm/model_executor/kernels/linear/ + commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt - label: LM Eval TurboQuant KV Cache key: lm-eval-turboquant-kv-cache diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 5227bbc1f3b..3a113f1982a 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -36,10 +36,10 @@ steps: source_file_dependencies: - vllm/ - tests/models/test_terratorch.py - - tests/models/test_transformers.py + - tests/models/transformers/test_backend.py - tests/models/test_registry.py commands: - - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py + - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py mirror: amd: device: mi325_1 @@ -55,6 +55,7 @@ steps: - vllm/ - tests/models/test_utils.py - tests/models/test_vision.py + - tests/models/transformers/fusers/ device: cpu-small commands: - - pytest -v -s models/test_utils.py models/test_vision.py + - pytest -v -s models/test_utils.py models/test_vision.py models/transformers/fusers/ diff --git a/.buildkite/test_areas/models_distributed.yaml b/.buildkite/test_areas/models_distributed.yaml index b5758c55aff..a3ee7666ed0 100644 --- a/.buildkite/test_areas/models_distributed.yaml +++ b/.buildkite/test_areas/models_distributed.yaml @@ -17,7 +17,7 @@ steps: - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)' # Avoid importing model tests that cause CUDA reinitialization error - - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)' + - pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)' - pytest models/language -v -s -m 'distributed(num_gpus=2)' - pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)' - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 8ca6fc22d64..57166d9d9b7 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -119,7 +119,7 @@ # Transformers modeling backend /vllm/model_executor/models/transformers @hmellor -/tests/models/test_transformers.py @hmellor +/tests/models/transformers @hmellor # Docs /docs/mkdocs @hmellor diff --git a/CMakeLists.txt b/CMakeLists.txt index 901f2be6bbb..48c0270e2c6 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -400,7 +400,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") "csrc/libtorch_stable/topk.cu" "csrc/libtorch_stable/mamba/selective_scan_fwd.cu" "csrc/libtorch_stable/cache_kernels.cu" - "csrc/libtorch_stable/cache_kernels.cu" "csrc/libtorch_stable/cache_kernels_fused.cu" "csrc/libtorch_stable/custom_all_reduce.cu" "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu") diff --git a/benchmarks/benchmark_topk_topp.py b/benchmarks/benchmark_topk_topp.py index 27b6dd8d6be..f727f16ea29 100644 --- a/benchmarks/benchmark_topk_topp.py +++ b/benchmarks/benchmark_topk_topp.py @@ -132,8 +132,10 @@ def benchmark_function( reset_memory_stats() # Benchmark - start_events = [torch.Event(enable_timing=True) for _ in range(benchmark_iters)] - end_events = [torch.Event(enable_timing=True) for _ in range(benchmark_iters)] + start_events = [ + torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters) + ] + end_events = [torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)] for i in range(benchmark_iters): logits_copy = logits.clone() diff --git a/benchmarks/kernels/benchmark_moe_defaults.py b/benchmarks/kernels/benchmark_moe_defaults.py index 7f000e01137..f6ad59366dc 100644 --- a/benchmarks/kernels/benchmark_moe_defaults.py +++ b/benchmarks/kernels/benchmark_moe_defaults.py @@ -134,8 +134,8 @@ def benchmark_config( torch.accelerator.synchronize() # Benchmark - start = torch.Event(enable_timing=True) - end = torch.Event(enable_timing=True) + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) start.record() for _ in range(num_iters): with override_config(config): diff --git a/benchmarks/kernels/benchmark_selective_state_update.py b/benchmarks/kernels/benchmark_selective_state_update.py index 5a3a6e88a63..a8b73da2aa9 100644 --- a/benchmarks/kernels/benchmark_selective_state_update.py +++ b/benchmarks/kernels/benchmark_selective_state_update.py @@ -170,8 +170,8 @@ def benchmark_config( graph.replay() torch.accelerator.synchronize() - start = torch.Event(enable_timing=True) - end = torch.Event(enable_timing=True) + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) latencies: list[float] = [] for _ in range(num_iters): start.record() diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 177c420776f..3aca9bcea91 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -15,6 +15,7 @@ endif() # set(ENABLE_X86_ISA $ENV{VLLM_CPU_X86}) set(ENABLE_ARM_BF16 $ENV{VLLM_CPU_ARM_BF16}) +set(ENABLE_RVV_BF16 $ENV{VLLM_CPU_RVV_BF16}) include_directories("${CMAKE_SOURCE_DIR}/csrc") @@ -110,6 +111,13 @@ else() set(ARM_BF16_FOUND ON) message(STATUS "ARM BF16 support enabled via VLLM_CPU_ARM_BF16 environment variable") endif() + # Some kernels (e.g. Bianbu on Spacemit X100) do not report zvfbfmin + # in /proc/cpuinfo despite hardware support. VLLM_CPU_RVV_BF16=1 + # overrides the detection result. + if (ENABLE_RVV_BF16) + set(RVV_BF16_FOUND ON) + message(STATUS "RVV BF16 support enabled via VLLM_CPU_RVV_BF16 environment variable") + endif() endif() if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64" OR ENABLE_X86_ISA) @@ -178,7 +186,10 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64") # Override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256 for RVV. if(NOT DEFINED VLLM_RVV_VLEN) # Auto-detect: find the largest zvlb in /proc/cpuinfo isa line. - if(EXISTS /proc/cpuinfo) + # Skip when cross-compiling — /proc/cpuinfo describes the build host. + if(CMAKE_CROSSCOMPILING) + message(STATUS "Cross-compiling: skipping VLEN auto-detection from /proc/cpuinfo") + elseif(EXISTS /proc/cpuinfo) file(READ /proc/cpuinfo _cpuinfo) set(_best 0) foreach(_n IN ITEMS 128 256 512 1024) @@ -186,6 +197,13 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64") set(_best ${_n}) endif() endforeach() + # Only VLEN=128 and VLEN=256 are supported by the RVV kernels. + if(_best GREATER 256) + message(WARNING + "Detected VLEN=${_best} but only 128/256 are supported; " + "clamping to 256") + set(_best 256) + endif() if(_best GREATER 0) set(VLLM_RVV_VLEN ${_best}) endif() @@ -195,9 +213,9 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64") if(NOT DEFINED VLLM_RVV_VLEN AND (RVV_FP16_FOUND OR RVV_BF16_FOUND)) message(FATAL_ERROR "RISC-V RVV is available but VLEN could not be auto-detected. " - "Please specify VLEN explicitly:\n" - " -DVLLM_RVV_VLEN=128 (for VLEN=128 hardware)\n" - " -DVLLM_RVV_VLEN=256 (for VLEN=256 hardware, e.g. Spacemit X100)") + "Please specify VLEN explicitly via CMAKE_ARGS:\n" + " CMAKE_ARGS='-DVLLM_RVV_VLEN=128' (for VLEN=128 hardware)\n" + " CMAKE_ARGS='-DVLLM_RVV_VLEN=256' (for VLEN=256 hardware, e.g. Spacemit X100)") endif() endif() if(VLLM_RVV_VLEN AND VLLM_RVV_VLEN GREATER 0) @@ -209,7 +227,7 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64") message(STATUS "BF16 extension detected") set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d) elseif(RVV_FP16_FOUND) - message(WARNING "BF16 functionality is not available") + message(WARNING "BF16 functionality is not available.") set(MARCH_FLAGS -march=rv64gcv_zvfh_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d) else() message(STATUS "compile riscv with scalar (no FP16/BF16)") diff --git a/csrc/cpu/cpu_attn.cpp b/csrc/cpu/cpu_attn.cpp index ec1a2b162de..fa22861157e 100644 --- a/csrc/cpu/cpu_attn.cpp +++ b/csrc/cpu/cpu_attn.cpp @@ -13,7 +13,8 @@ static inline cpu_attention::Fp8KVCacheDataType parse_fp8_kv_dtype( bool cpu_attn_has_isa(const std::string& isa) { if (isa == "rvv") { -#if defined(__riscv) && defined(__riscv_v_min_vlen) && __riscv_v_min_vlen == 128 +#if defined(__riscv) && defined(__riscv_v_min_vlen) && \ + (__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256) return true; #else return false; diff --git a/csrc/cpu/cpu_types_riscv_impl.hpp b/csrc/cpu/cpu_types_riscv_impl.hpp index d0ce67a5afe..70cb0ab52de 100644 --- a/csrc/cpu/cpu_types_riscv_impl.hpp +++ b/csrc/cpu/cpu_types_riscv_impl.hpp @@ -214,11 +214,18 @@ struct BF16Vec32 : public Vec { explicit BF16Vec32(const BF16Vec8& v) { fixed_u16x8_t u16_val = bf16_to_u16(v.reg); - fixed_u16x32_t u16_combined = - RVVI4(__riscv_vcreate_v_u16, LMUL_128, _u16, LMUL_512)( - u16_val, u16_val, u16_val, u16_val); - reg = RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16, - LMUL_512)(u16_combined); + // Widen LMUL_128 → LMUL_256 so vslideup operands share a type. + // At VLEN=256 this is mf2→m1 (both integer); at VLEN=128 it is m1→m2. + fixed_u16x16_t ext = + RVVI4(__riscv_vlmul_ext_v_u16, LMUL_128, _u16, LMUL_256)(u16_val); + // Build 16-element half: place the 8 elements at offsets 0 and 8. + fixed_u16x16_t half = RVVI(__riscv_vmv_v_x_u16, LMUL_256)(0, 16); + half = RVVI(__riscv_vslideup_vx_u16, LMUL_256)(half, ext, 0, 8); + half = RVVI(__riscv_vslideup_vx_u16, LMUL_256)(half, ext, 8, 16); + // Double to LMUL_512 (m1→m2 at VLEN=256, m2→m4 at VLEN=128). + fixed_u16x32_t dst = + RVVI4(__riscv_vcreate_v_u16, LMUL_256, _u16, LMUL_512)(half, half); + reg = RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16, LMUL_512)(dst); }; void save(void* ptr) const { @@ -623,17 +630,29 @@ struct FP32Vec16 : public Vec { data.reg, data.reg)) {}; explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {}; explicit FP32Vec16(int64_t value, const FP32Vec16& lut) { - const uint64_t q_values = static_cast(value); - auto packed = RVVI(__riscv_vmv_v_x_u64, LMUL_1024)(q_values, VEC_ELEM_NUM); - auto lane_ids = RVVI(__riscv_vid_v_u64, LMUL_1024)(VEC_ELEM_NUM); - auto shifts = - RVVI(__riscv_vsll_vx_u64, LMUL_1024)(lane_ids, 2, VEC_ELEM_NUM); - auto shifted = - RVVI(__riscv_vsrl_vv_u64, LMUL_1024)(packed, shifts, VEC_ELEM_NUM); - auto idx64 = - RVVI(__riscv_vand_vx_u64, LMUL_1024)(shifted, 0xF, VEC_ELEM_NUM); - auto idx32 = RVVI(__riscv_vnsrl_wx_u32, LMUL_512)(idx64, 0, VEC_ELEM_NUM); - reg = RVVI(__riscv_vrgather_vv_f32, LMUL_512)(lut.reg, idx32, VEC_ELEM_NUM); + // Split into two 32-bit halves to avoid u64 @ LMUL_1024 (m8 on + // VLEN=128 / m4 on VLEN=256), which causes heavy register spilling. + constexpr int HALF = VEC_ELEM_NUM / 2; + const auto q = static_cast(value); + const uint32_t lo = static_cast(q); + const uint32_t hi = static_cast(q >> 32); + + auto lane_ids = RVVI(__riscv_vid_v_u32, LMUL_256)(HALF); + auto shifts = RVVI(__riscv_vsll_vx_u32, LMUL_256)(lane_ids, 2, HALF); + + auto packed_lo = RVVI(__riscv_vmv_v_x_u32, LMUL_256)(lo, HALF); + auto idx_lo = RVVI(__riscv_vand_vx_u32, LMUL_256)( + RVVI(__riscv_vsrl_vv_u32, LMUL_256)(packed_lo, shifts, HALF), 0xF, + HALF); + + auto packed_hi = RVVI(__riscv_vmv_v_x_u32, LMUL_256)(hi, HALF); + auto idx_hi = RVVI(__riscv_vand_vx_u32, LMUL_256)( + RVVI(__riscv_vsrl_vv_u32, LMUL_256)(packed_hi, shifts, HALF), 0xF, + HALF); + + auto idx = + RVVI4(__riscv_vcreate_v_u32, LMUL_256, _u32, LMUL_512)(idx_lo, idx_hi); + reg = RVVI(__riscv_vrgather_vv_f32, LMUL_512)(lut.reg, idx, VEC_ELEM_NUM); } explicit FP32Vec16(const FP16Vec16& v); diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index e17c9ab3a7e..cfa296e73b6 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -278,7 +278,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.def( "dynamic_4bit_int_moe(" "Tensor x, Tensor topk_ids, Tensor topk_weights," - "Tensor w13_packed, Tensor w2_packed, int H, int I, int I2," + "Tensor w13_packed, Tensor w2_packed," + "int hidden_size, int intermediate_size," "int group_size, bool apply_router_weight_on_input, int activation_kind" ") -> Tensor"); diff --git a/csrc/moe/dynamic_4bit_int_moe_cpu.cpp b/csrc/moe/dynamic_4bit_int_moe_cpu.cpp index 58dc4020168..1b071d334ff 100644 --- a/csrc/moe/dynamic_4bit_int_moe_cpu.cpp +++ b/csrc/moe/dynamic_4bit_int_moe_cpu.cpp @@ -29,25 +29,37 @@ enum ActivationKind : int64_t { torch::Tensor dynamic_4bit_int_moe_cpu( torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights, - torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t H, int64_t I, - int64_t I2, int64_t group_size, bool apply_router_weight_on_input, - int64_t activation_kind) { + torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t hidden_size, + int64_t intermediate_size, int64_t group_size, + bool apply_router_weight_on_input, int64_t activation_kind) { TORCH_CHECK(x.dim() == 2, "x must be 2D"); TORCH_CHECK(topk_ids.dim() == 2 && topk_weights.dim() == 2, "topk tensors must be [T, K]"); TORCH_CHECK( w13_packed.size(0) == w2_packed.size(0), "w13_packed and w2_packed must have same number of experts in dim 0"); - TORCH_CHECK(I2 == 2 * I, "I2 must equal 2*I"); const int64_t T = x.size(0); const int64_t K = topk_ids.size(1); const int64_t E = w13_packed.size(0); const int64_t N = T * K; + const int64_t w13_out_features = 2 * intermediate_size; auto x_c = x.contiguous(); + // _dyn_quant_matmul_4bit kernel natively supports these pre-quant activation + // dtypes: + // - fp32: with channelwise and groupwise + // - bf16: with channelwise -> upcast to fp32 for groupwise + // - fp16: not supported -> upcast to fp32 for groupwise & channelwise + const auto output_dtype = x_c.scalar_type(); + const bool should_cast_input = + ((group_size != -1) && output_dtype == at::kBFloat16) || + output_dtype == at::kHalf; + if (should_cast_input) { + x_c = x_c.to(at::kFloat); + } auto ids_c = topk_ids.contiguous(); - auto gates_c = topk_weights.to(at::kFloat).contiguous(); + auto gates_c = topk_weights.to(x_c.scalar_type()).contiguous(); // bucketing tokens -> experts c10::SmallVector counts( @@ -63,35 +75,42 @@ torch::Tensor dynamic_4bit_int_moe_cpu( c10::SmallVector offsets(E + 1, 0); // ( E +1 ) for (int64_t e = 0; e < E; ++e) offsets[e + 1] = offsets[e] + counts[e]; + // expert_tokens = [tokens indices for expert 0, ...] + // expert_gates = [router weights for tokens assigned to expert 0, ...] auto expert_tokens = at::empty({offsets[E]}, ids_c.options()); auto expert_gates = at::empty({offsets[E]}, gates_c.options()); { c10::SmallVector cursor(E, 0); - const auto* ids_ptr = ids_c.data_ptr(); - const auto* gts_ptr = gates_c.data_ptr(); - auto* tok_ptr = expert_tokens.data_ptr(); - auto* gate_ptr = expert_gates.data_ptr(); + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::BFloat16, at::ScalarType::Half, gates_c.scalar_type(), + "bucket_expert_tokens_and_gates", [&] { + const auto* ids_ptr = ids_c.data_ptr(); + const auto* gts_ptr = gates_c.data_ptr(); + auto* tok_ptr = expert_tokens.data_ptr(); + auto* gate_ptr = expert_gates.data_ptr(); - for (int64_t t = 0; t < T; ++t) { - const int64_t base = t * K; - for (int64_t k = 0; k < K; ++k) { - const int64_t idx = base + k; - const int64_t e = ids_ptr[idx]; - const int64_t p = offsets[e] + (cursor[e]++); - tok_ptr[p] = t; - gate_ptr[p] = gts_ptr[idx]; - } - } + for (int64_t t = 0; t < T; ++t) { + const int64_t base = t * K; + for (int64_t k = 0; k < K; ++k) { + const int64_t idx = base + k; + const int64_t e = ids_ptr[idx]; + const int64_t p = offsets[e] + (cursor[e]++); + tok_ptr[p] = t; + gate_ptr[p] = gts_ptr[idx]; + } + } + }); } - const int64_t g_eff_13 = (group_size != -1) ? group_size : H; - const int64_t g_eff_2 = (group_size != -1) ? group_size : I; + const int64_t g_eff_13 = (group_size != -1) ? group_size : hidden_size; + const int64_t g_eff_2 = (group_size != -1) ? group_size : intermediate_size; + // X_all [num_tokens * K, hidden_size] auto X_all = x_c.index_select(/*dim=*/0, expert_tokens); if (apply_router_weight_on_input) { X_all = X_all.mul(expert_gates.unsqueeze(1)); } - auto Y_all = at::empty({offsets[E], H}, x_c.options()); + auto Y_all = at::empty({offsets[E], hidden_size}, x_c.options()); at::parallel_for(0, offsets[E], 0, [&](int64_t idx_begin, int64_t idx_end) { c10::InferenceMode guard; @@ -109,11 +128,13 @@ torch::Tensor dynamic_4bit_int_moe_cpu( auto w2_e = w2_packed.select(/*dim=*/0, e); // W13 - auto y13 = - mm(x_e, w13_e, g_eff_13, /*in_features=*/H, /*out_features=*/I2); + auto y13 = mm(x_e, w13_e, g_eff_13, /*in_features=*/hidden_size, + /*out_features=*/w13_out_features); - auto g_part = y13.narrow(/*dim=*/1, /*start=*/0, /*length=*/I); - auto u_part = y13.narrow(/*dim=*/1, /*start=*/I, /*length=*/I); + auto g_part = + y13.narrow(/*dim=*/1, /*start=*/0, /*length=*/intermediate_size); + auto u_part = y13.narrow(/*dim=*/1, /*start=*/intermediate_size, + /*length=*/intermediate_size); torch::Tensor act; if (activation_kind == ActivationKind::SwiGLUOAI) { // SwiGLUOAI @@ -128,7 +149,8 @@ torch::Tensor dynamic_4bit_int_moe_cpu( } // W2 - auto y = mm(act, w2_e, g_eff_2, /*in_features=*/I, /*out_features=*/H); + auto y = mm(act, w2_e, g_eff_2, /*in_features=*/intermediate_size, + /*out_features=*/hidden_size); // Store per-expert result Y_all.narrow(/*dim=*/0, /*start=*/start, /*length=*/te).copy_(y); @@ -138,8 +160,11 @@ torch::Tensor dynamic_4bit_int_moe_cpu( if (!apply_router_weight_on_input) { Y_all = Y_all.mul(expert_gates.unsqueeze(1)); } + if (Y_all.scalar_type() != output_dtype) { + Y_all = Y_all.to(output_dtype); + } - auto out = at::zeros({T, H}, x.options()); + auto out = at::zeros({T, hidden_size}, x.options()); out = at::index_add(out, /*dim=*/0, /*index=*/expert_tokens, /*source=*/Y_all); diff --git a/csrc/ops.h b/csrc/ops.h index cd18b1e5e0d..274cd52bea4 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -53,9 +53,9 @@ void dynamic_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input, torch::Tensor dynamic_4bit_int_moe_cpu( torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights, - torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t H, int64_t I, - int64_t I2, int64_t group_size, bool apply_router_weight_on_input, - int64_t activation_kind); + torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t hidden_size, + int64_t intermediate_size, int64_t group_size, + bool apply_router_weight_on_input, int64_t activation_kind); using fptr_t = int64_t; #ifdef USE_ROCM diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu index ee5d5daf649..a528ffbd8d1 100644 --- a/docker/Dockerfile.cpu +++ b/docker/Dockerfile.cpu @@ -25,7 +25,6 @@ FROM ubuntu:22.04 AS base-common WORKDIR /workspace ARG PYTHON_VERSION=3.12 -ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cpu" ARG max_jobs=32 ENV MAX_JOBS=${max_jobs} @@ -53,8 +52,6 @@ ENV PATH="$VIRTUAL_ENV/bin:$PATH" ENV UV_HTTP_TIMEOUT=500 # Install Python dependencies -ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL} -ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL} ENV UV_INDEX_STRATEGY="unsafe-best-match" ENV UV_LINK_MODE="copy" @@ -64,7 +61,7 @@ COPY requirements/cpu.txt requirements/cpu.txt RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install --upgrade pip && \ - uv pip install -r requirements/cpu.txt + uv pip install -r requirements/cpu.txt --torch-backend cpu ARG TARGETARCH ENV TARGETARCH=${TARGETARCH} @@ -149,7 +146,7 @@ RUN if [ "$TARGETARCH" = "arm64" ] && [ "$VLLM_CPU_X86" != "0" ]; then \ COPY requirements/build/cpu.txt requirements/build/cpu.txt RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install -r requirements/build/cpu.txt + uv pip install -r requirements/build/cpu.txt --torch-backend cpu COPY . . @@ -205,7 +202,7 @@ RUN case "$(uname -m)" in \ esac RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install -r requirements/test/cpu.txt + uv pip install -r requirements/test/cpu.txt --torch-backend cpu ######################### DEV IMAGE ######################### FROM vllm-build AS vllm-dev @@ -231,7 +228,7 @@ COPY --from=vllm-test-deps /vllm-workspace/requirements/test/cpu.txt requirement RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install -r requirements/lint.txt && \ - uv pip install -r requirements/test/cpu.txt && \ + uv pip install -r requirements/test/cpu.txt --torch-backend cpu && \ pre-commit install --hook-type pre-commit --hook-type commit-msg ENTRYPOINT ["bash"] diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 3f307a5fa0f..02e4086d625 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -252,6 +252,8 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/docker-bake.hcl /docker/doc COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/docker-bake-rocm.hcl /docker/docker-bake-rocm.hcl COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust /rust +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 # RIXL/UCX build stages @@ -543,6 +545,8 @@ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/docker-bake.h COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/docker-bake-rocm.hcl /docker/docker-bake-rocm.hcl COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/rust /rust +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 # ----------------------- @@ -576,6 +580,7 @@ RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ libibverbs1 \ ibverbs-providers \ ibverbs-utils \ + unzip \ pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \ libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \ && rm -rf /var/lib/apt/lists/* diff --git a/docker/Dockerfile.s390x b/docker/Dockerfile.s390x index 554a7257c23..6d1c0c3452f 100644 --- a/docker/Dockerfile.s390x +++ b/docker/Dockerfile.s390x @@ -249,7 +249,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ NUMBA_WHL_FILE=$(ls /tmp/numba-wheels/*.whl) && \ OPENCV_WHL_FILE=$(ls /tmp/opencv-wheels/*.whl) && \ GUIDANCE_WHL_FILE=$(ls /tmp/guidance-wheels/*.whl) && \ - uv pip install -v \ + uv pip install -v \ $ARROW_WHL_FILE \ $VISION_WHL_FILE \ $HF_XET_WHL_FILE \ @@ -257,6 +257,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ $NUMBA_WHL_FILE \ $OPENCV_WHL_FILE \ $GUIDANCE_WHL_FILE \ + --torch-backend cpu \ --index-strategy unsafe-best-match \ -r requirements/build/cpu.txt \ -r requirements/cpu.txt diff --git a/docs/configuration/optimization.md b/docs/configuration/optimization.md index c6d64b25035..0f407fbc97f 100644 --- a/docs/configuration/optimization.md +++ b/docs/configuration/optimization.md @@ -276,8 +276,9 @@ By default vLLM uses the standard Hugging Face `tokenizers` library to power the fast tokenizer. For BPE tokenizers (Qwen, Llama, DeepSeek, GPT-OSS, etc.) you can switch to the [fastokens](https://github.com/crusoecloud/fastokens) Rust backend, a drop-in replacement that's substantially faster on -encode/decode and on streaming detokenization. Enable it by setting -`VLLM_USE_FASTOKENS=1`: +encode/decode and on streaming detokenization. `VLLM_USE_FASTOKENS` is +available in vLLM v0.23.0 and later. If your installed vLLM version does not +recognize the environment variable, upgrade vLLM before enabling the override: ```console VLLM_USE_FASTOKENS=1 vllm serve Qwen/Qwen3-8B diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index d28de8fc36a..a252b59e9e4 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -168,7 +168,7 @@ Priority is **1 = highest** (tried first). | `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 | | `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any | | `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any | -| `HPC_ATTN` | | fp16, bf16 | `auto`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 | +| `HPC_ATTN` | | fp16, bf16 | `auto`, `bfloat16`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 | | `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A | | `ROCM_AITER_UNIFIED_ATTN` | | bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A | | `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A | diff --git a/docs/features/multimodal_inputs.md b/docs/features/multimodal_inputs.md index 847743dfff1..e44596626f8 100644 --- a/docs/features/multimodal_inputs.md +++ b/docs/features/multimodal_inputs.md @@ -816,6 +816,44 @@ Full example: [examples/generate/multimodal/openai_chat_completion_client_for_mu export VLLM_VIDEO_FETCH_TIMEOUT= ``` +#### Video Decoding Backend + +vLLM decodes video bytes into frames using a selectable decoding backend. Three +backends are supported: + +- `opencv` (default): OpenCV-based decoder. +- `pyav`: PyAV decoder. +- `torchcodec`: TorchCodec (PyTorch-native) decoder. + +All three backends are ultimately backed by FFmpeg. `torchcodec` lets +you choose which FFmpeg version is used while `opencv` and `pyav` rely on +whichever FFmpeg build they were linked against. + +Select the backend by passing the `backend` parameter via `--media-io-kwargs`: + +```bash +vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \ + --media-io-kwargs '{"video": {"backend": "torchcodec"}}' +``` + +**TorchCodec-specific parameters:** + +The following parameters only apply to the `torchcodec` backend: + +- `num_ffmpeg_threads`: Number of FFmpeg decoding threads. `0` (default) relies + on the FFmpeg default, which is `min(cpu_count + 1, 16)`. This allows you to + control thread over-subscription. +- `seek_mode`: Seek mode for the decoder. `"exact"` (default) guarantees + frame-accurate sampling by scanning the file when the decoder is created. + `"approximate"` skips that scan for faster decoder creation, at the cost of + relying on the file's metadata (which may yield less accurate seeking). + +```bash +# Example: TorchCodec with approximate seek mode and 4 FFmpeg threads +vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \ + --media-io-kwargs '{"video": {"backend": "torchcodec", "seek_mode": "approximate", "num_ffmpeg_threads": 4}}' +``` + #### Video Frame Recovery For improved robustness when processing potentially corrupted or truncated video files, vLLM supports optional frame recovery using a dynamic window forward-scan approach. When enabled, if a target frame fails to load during sequential reading, the next successfully grabbed frame (before the next target frame) will be used in its place. diff --git a/docs/features/per_request_metrics.md b/docs/features/per_request_metrics.md new file mode 100644 index 00000000000..9bc64d2b86d --- /dev/null +++ b/docs/features/per_request_metrics.md @@ -0,0 +1,127 @@ +# Per-Request Metrics + +vLLM can return per-request timing metrics directly in API responses. +This is useful for billing, SLA monitoring, and latency analysis at the +individual request level, as a complement to the server-aggregated Prometheus +metrics exposed at `/metrics`. + +## Enabling + +Start the server with `--enable-per-request-metrics`: + +```bash +vllm serve meta-llama/Llama-3.1-8B-Instruct --enable-per-request-metrics +``` + +When this flag is set, supported API responses include metrics for each +attributable request. + +!!! note + At high concurrency, enabling per-request metrics computation may introduce + non-negligible CPU overhead. Benchmark your specific workload to evaluate the + impact before enabling in production. + +## Response Format + +When per-request metrics are enabled, the response includes a `metrics` object: + +```json +{ + "id": "chatcmpl-abc123", + "object": "chat.completion", + "model": "meta-llama/Llama-3.1-8B-Instruct", + "choices": [ ... ], + "usage": { + "prompt_tokens": 42, + "completion_tokens": 128, + "total_tokens": 170 + }, + "metrics": { + "time_to_first_token_ms": 85.2, + "generation_time_ms": 1240.5, + "queue_time_ms": 12.3, + "mean_itl_ms": 9.1, + "tokens_per_second": 103.2 + } +} +``` + +| Field | Description | +| --- | --- | +| `time_to_first_token_ms` | Time from when the request was scheduled until the first output token was generated (TTFT). | +| `generation_time_ms` | Decode time: time from the first output token to the last output token. Excludes both queue wait and prefill/TTFT. | +| `queue_time_ms` | Time the request spent waiting in the scheduler queue before processing began. | +| `mean_itl_ms` | Mean inter-token latency (average time between successive output tokens) during the decode phase. `null` for single-token responses. | +| `tokens_per_second` | Overall output token throughput: all generated tokens over the inference interval (scheduling to last output token). Unlike `generation_time_ms`, this includes the prefill phase, so it reflects end-to-end generation speed rather than pure decode speed. | + +All fields are `null` if the underlying timing data is not available for that +request. + +!!! note + Timing metrics describe a single generation stream, so they are only + returned when the request maps to exactly one. They are suppressed (the + `metrics` object is `null`) for requests with `n > 1`, because the + underlying timing data reflects only one of the `n` sequences and cannot be + accurately attributed to the request as a whole. Token usage + (`prompt_tokens`, `completion_tokens`) remains accurate in these cases. + Per-request metrics also require server-side statistics logging, which is + on by default. vLLM rejects `--enable-per-request-metrics` when + `--disable-log-stats` is also set. + +## Example Request + +=== "Non-streaming" + + ```python + from openai import OpenAI + + client = OpenAI(base_url="http://localhost:8000/v1", api_key="token") + + response = client.chat.completions.create( + model="meta-llama/Llama-3.1-8B-Instruct", + messages=[{"role": "user", "content": "What is the capital of France?"}], + ) + + print(response.usage) + print(response.model_extra.get("metrics")) + ``` + +=== "Streaming" + + In streaming responses, metrics are attached to the final usage chunk (the + chunk sent after all content chunks). That chunk is only emitted when usage + reporting is enabled with `stream_options.include_usage: true` or forced + server-side with `--enable-force-include-usage`. Without forced usage, a + streaming client must set `stream_options.include_usage: true` to receive + metrics. + + ```python + from openai import OpenAI + + client = OpenAI(base_url="http://localhost:8000/v1", api_key="token") + + stream = client.chat.completions.create( + model="meta-llama/Llama-3.1-8B-Instruct", + messages=[{"role": "user", "content": "What is the capital of France?"}], + stream=True, + stream_options={"include_usage": True}, + ) + + for chunk in stream: + if chunk.usage: + print("Usage:", chunk.usage) + print("Metrics:", chunk.model_extra.get("metrics")) + ``` + +## Completions API + +Per-request metrics are also available on the `/v1/completions` endpoint using +the same `metrics` response field. As with `n > 1`, metrics are omitted for +requests with multiple prompts, because the timing data cannot be attributed to +a single prompt's generation. + +## Relationship to Prometheus Metrics + +The `metrics` response field provides per-request values for a single request. +The `/metrics` Prometheus endpoint exposes server-level histograms (e.g. +`vllm:time_to_first_token_seconds`) that aggregate across all requests. diff --git a/docs/features/quantization/inc.md b/docs/features/quantization/inc.md index adb6b3ae8e2..ffb90cec8c1 100644 --- a/docs/features/quantization/inc.md +++ b/docs/features/quantization/inc.md @@ -75,14 +75,11 @@ vllm serve Intel/DeepSeek-R1-0528-Qwen3-8B-int4-AutoRound \ --max-model-len 4096 ``` -!!! note - To deploy `wNa16` models on Intel GPU/CPU, please add `--enforce-eager` for now. - ## Evaluating the Quantized Model with vLLM ```bash lm_eval --model vllm \ - --model_args pretrained="Intel/DeepSeek-R1-0528-Qwen3-8B-int4-AutoRound,max_model_len=8192,max_num_batched_tokens=32768,max_num_seqs=128,gpu_memory_utilization=0.8,dtype=bfloat16,max_gen_toks=2048,enforce_eager=True" \ + --model_args pretrained="Intel/DeepSeek-R1-0528-Qwen3-8B-int4-AutoRound,max_model_len=8192,max_num_batched_tokens=32768,max_num_seqs=128,gpu_memory_utilization=0.8,dtype=bfloat16,max_gen_toks=2048" \ --tasks gsm8k \ --num_fewshot 5 \ --batch_size 128 diff --git a/docs/getting_started/installation/cpu.apple.inc.md b/docs/getting_started/installation/cpu.apple.inc.md index e312964ec8a..479b6d2c011 100644 --- a/docs/getting_started/installation/cpu.apple.inc.md +++ b/docs/getting_started/installation/cpu.apple.inc.md @@ -35,15 +35,10 @@ After installation of XCode and the Command Line Tools, which include Apple Clan ```bash git clone https://github.com/vllm-project/vllm.git cd vllm -uv pip install -r requirements/cpu.txt --index-strategy unsafe-best-match +uv pip install -r requirements/cpu.txt uv pip install -e . ``` -!!! tip - The `--index-strategy unsafe-best-match` flag is needed to resolve dependencies across multiple package indexes (PyTorch CPU index and PyPI). Without this flag, you may encounter `typing-extensions` version conflicts. - - The term "unsafe" refers to the package resolution strategy, not security. By default, `uv` only searches the first index where a package is found to prevent dependency confusion attacks. This flag allows `uv` to search all configured indexes to find the best compatible versions. Since both PyTorch and PyPI are trusted package sources, using this strategy is safe and appropriate for vLLM installation. - !!! note On macOS the `VLLM_TARGET_DEVICE` is automatically set to `cpu`, which is currently the only supported device. diff --git a/docs/getting_started/installation/cpu.s390x.inc.md b/docs/getting_started/installation/cpu.s390x.inc.md index 1e36b431764..15baa487c2a 100644 --- a/docs/getting_started/installation/cpu.s390x.inc.md +++ b/docs/getting_started/installation/cpu.s390x.inc.md @@ -48,10 +48,10 @@ Execute the following commands to build and install vLLM from source. ```bash uv pip install -v \ - --extra-index-url https://download.pytorch.org/whl/cpu \ - --torch-backend auto \ -r requirements/build/cpu.txt \ -r requirements/cpu.txt \ + --torch-backend cpu \ + --index-strategy unsafe-best-match && \ VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \ uv pip install dist/*.whl ``` diff --git a/docs/models/pooling_models/README.md b/docs/models/pooling_models/README.md index 37fca366eba..f8de9d437ad 100644 --- a/docs/models/pooling_models/README.md +++ b/docs/models/pooling_models/README.md @@ -2,8 +2,7 @@ !!! note We currently support pooling models primarily for convenience. This is not guaranteed to provide any performance -improvements over using Hugging Face Transformers or Sentence Transformers directly. - + improvements over using Hugging Face Transformers or Sentence Transformers directly. We plan to optimize pooling models in vLLM. Please comment on if you have any suggestions! ## What are pooling models? @@ -63,7 +62,7 @@ please refer to [IO Processor Plugins](../../design/io_processor_plugins.md). !!! note Within classification tasks, there is a specialized subcategory: Cross-encoder (aka reranker) models. These models -are a subset of classification models that accept two prompts as input and output num_labels equal to 1. + are a subset of classification models that accept two prompts as input and output num_labels equal to 1. ### Pooling Types diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 562e38109ff..6b239b6f2fd 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -15,7 +15,7 @@ These models are what we list in [supported text models](#list-of-text-only-lang ### Transformers -vLLM also supports model implementations that are available in Transformers. You should expect the performance of a Transformers model implementation used in vLLM to be within <5% of the performance of a dedicated vLLM model implementation. We call this feature the "Transformers modeling backend". +vLLM also supports model implementations that are available in Transformers. We call this feature the "Transformers modeling backend". The performance of models loaded with the Transformers modeling backend should be identical to a dedicated vLLM model implementation. Currently, the Transformers modeling backend works for the following: @@ -140,7 +140,7 @@ Here is what happens in the background when this model is loaded: That's it! -For your model to be compatible with vLLM's tensor parallel and/or pipeline parallel features, you must add `base_model_tp_plan` and/or `base_model_pp_plan` to your model's config class: +For your model to be compatible with vLLM's tensor parallel and/or pipeline parallel features, you may need to add `base_model_tp_plan` and/or `base_model_pp_plan` to your model's config class:
configuration_my_model.py @@ -168,9 +168,11 @@ class MyConfig(PretrainedConfig):
- `base_model_tp_plan` is a `dict` that maps fully qualified layer name patterns to tensor parallel styles (currently only `"colwise"` and `"rowwise"` are supported). + - vLLM infers the tensor parallel style of standard attention (`q`/`k`/`v`/`o_proj`) and gated-MLP/experts (`gate`/`up`/`down_proj`) projections if it can fuse them, so these may not need to be listed. `base_model_tp_plan` is only _required_ for layers that do not follow these patterns; any linear that is neither fused nor named in the plan is replicated. - `base_model_pp_plan` is a `dict` that maps direct child layer names to `tuple`s of `list`s of `str`s: - You only need to do this for layers which are not present on all pipeline stages - vLLM assumes that there will be only one `nn.ModuleList`, which is distributed across the pipeline stages + - When no `base_model_pp_plan` is provided, the Transformers modelling backend infers the split from the text model's sole `nn.ModuleList`, keeping the parameter-bearing modules around it (input embeddings, final norm) on the first/last stage (depending on declaration order) and parameter-free modules (e.g. rotary embeddings) on every stage - The `list` in the first element of the `tuple` contains the names of the input arguments - The `list` in the last element of the `tuple` contains the names of the variables the layer outputs to in your modeling code diff --git a/requirements/build/cpu.txt b/requirements/build/cpu.txt index 640432ddd8c..27a3ac65c98 100644 --- a/requirements/build/cpu.txt +++ b/requirements/build/cpu.txt @@ -1,4 +1,3 @@ ---extra-index-url https://download.pytorch.org/whl/cpu cmake>=3.26.1 ninja packaging>=24.2 diff --git a/requirements/cpu.txt b/requirements/cpu.txt index 5ec338af736..c0b98d22c9b 100644 --- a/requirements/cpu.txt +++ b/requirements/cpu.txt @@ -1,4 +1,3 @@ ---extra-index-url https://download.pytorch.org/whl/cpu # Common dependencies -r common.txt @@ -16,6 +15,9 @@ torchaudio; platform_machine != "s390x" and platform_machine != "riscv64" # required for the image processor of phi3v, this must be updated alongside torch torchvision; platform_machine != "s390x" and platform_machine != "riscv64" +# required for the torchcodec video decoding backend +torchcodec >= 0.14; platform_machine != "s390x" and platform_machine != "riscv64" and platform_machine != "ppc64le" + # Intel Extension for PyTorch, only for x86_64 CPUs intel-openmp==2024.2.1; platform_machine == "x86_64" diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 5545d3344f0..1d90c7ef404 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -8,6 +8,7 @@ torch==2.11.0 torchaudio==2.11.0 # These must be updated alongside torch torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version +torchcodec >= 0.14 PyNvVideoCodec==2.1.0 # FlashInfer should be updated together with the Dockerfile flashinfer-python==0.6.13 @@ -25,7 +26,7 @@ nvidia-cutlass-dsl[cu13]==4.5.2 quack-kernels>=0.3.3 # Tokenspeed_MLA for faster mla with spec decode -tokenspeed-mla==0.1.2 +tokenspeed-mla==0.1.2; platform_system == "Linux" # Humming kernels for quantization gemm -humming-kernels[cu13]==0.1.6 +humming-kernels[cu13]==0.1.10 diff --git a/requirements/test/cpu.txt b/requirements/test/cpu.txt index 6cc3a3782b8..5ae281c21df 100644 --- a/requirements/test/cpu.txt +++ b/requirements/test/cpu.txt @@ -1134,6 +1134,8 @@ torchaudio==2.11.0+cpu # -r requirements/test/cuda.in # encodec # vocos +torchcodec==0.14.0+cpu + # via -r requirements/test/cuda.in torchvision==0.26.0+cpu # via # -r requirements/test/cuda.in diff --git a/requirements/test/cuda.in b/requirements/test/cuda.in index 9a6e46712cb..c5518ecaa47 100644 --- a/requirements/test/cuda.in +++ b/requirements/test/cuda.in @@ -13,6 +13,7 @@ pytest-cov # testing utils albumentations # required for Nemotron Parse in test_common.py av # required for audio_in_video tests +torchcodec >= 0.14 # required for torchcodec video backend tests backoff # required for phi4mm test blobfile # required for kimi-vl test httpx diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index 3f7e4a5b5f3..53aaa78f95e 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -1233,6 +1233,10 @@ torchaudio==2.11.0+cu130 # -r requirements/test/cuda.in # encodec # vocos +torchcodec==0.14.0+cu130 + # via + # -c requirements/cuda.txt + # -r requirements/test/cuda.in torchvision==0.26.0+cu130 # via # -c requirements/cuda.txt diff --git a/requirements/xpu.txt b/requirements/xpu.txt index 684d0ef30f0..68b8eb13010 100644 --- a/requirements/xpu.txt +++ b/requirements/xpu.txt @@ -15,6 +15,7 @@ numba == 0.65.0 # Required for N-gram speculative decoding torch==2.12.0 torchaudio torchvision +torchcodec >= 0.14 # Required for the torchcodec video decoding backend -auto_round_lib>=0.13.3 +auto_round_lib>=0.14.0 vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.10.1/vllm_xpu_kernels-0.1.10.1-cp38-abi3-manylinux_2_28_x86_64.whl diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index 8284ddd1285..8e7ed02a403 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -172,6 +172,9 @@ impl ChatLlm { pub async fn chat(&self, mut request: ChatRequest) -> Result { request.validate()?; + // Stamp before rendering so render and tokenize count toward TTFT/e2e. + let arrival_time = vllm_llm::current_unix_timestamp_secs(); + let output_processor = self.backend.new_chat_output_processor( &mut request, NewChatOutputProcessorOptions { @@ -210,6 +213,7 @@ impl ChatLlm { data_parallel_rank: request.data_parallel_rank, reasoning_parser_kwargs, lora_request: request.lora_request, + arrival_time: Some(arrival_time), }; let decoded_stream = self.text.generate(text_request).await?.map_err(Error::from).boxed(); diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index e6044614aa9..b1670814ef6 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -130,6 +130,19 @@ impl RoundtripCase { } } + /// DeepSeek V3.2 DSML tool-call format. + fn deepseek_v32() -> Self { + Self { + model_id: "deepseek-ai/DeepSeek-V3.2-Exp", + assistant_stop_suffix: "<|end▁of▁sentence|>", + tool_call_parser: ParserSelection::Auto, + reasoning_parser: ParserSelection::Auto, + thinking_behavior: ThinkingBehavior::Toggleable { default: false }, + json_fmt: compact_json_fmt(), + sort_json_keys: false, + } + } + /// GLM-4.7 XML-like argument format with `` reasoning tags. fn glm47() -> Self { Self { @@ -235,6 +248,7 @@ roundtrip_tests! { qwen35 => [reasoning_and_content, tool_call_mix], minimax_m25 => [reasoning_and_content, tool_call_mix], deepseek_v4 => [reasoning_and_content, tool_call_mix], + deepseek_v32 => [tool_call_mix], glm47 => [reasoning_and_content, tool_call_mix], seed_oss => [reasoning_and_content], step3p5 => [reasoning_and_content], diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index babc8179329..eb899e3e784 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -512,6 +512,7 @@ impl EngineCoreClient { Ok(EngineCoreOutputStream::new( request_id, + engine_id.engine_index().unwrap_or(0), self.abort_tx.clone(), rx, )) diff --git a/rust/src/engine-core-client/src/client/imp.rs b/rust/src/engine-core-client/src/client/imp.rs index 3465817a453..c91a93d2714 100644 --- a/rust/src/engine-core-client/src/client/imp.rs +++ b/rust/src/engine-core-client/src/client/imp.rs @@ -15,7 +15,7 @@ use crate::client::state::{OutputReceiver, RequestRegistry, UtilityReceiver, Uti use crate::client::stream::EngineCoreStreamOutput; use crate::client::{AbortCause, AbortRequest}; use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_output}; -use crate::metrics::{LoraInfoExporter, record_scheduler_stats}; +use crate::metrics::{LoraInfoExporter, SchedulerStatsRecorder}; use crate::protocol::encode_msgpack; use crate::protocol::output::{EngineCoreOutput, EngineCoreOutputs}; use crate::protocol::request::EngineCoreRequestType; @@ -29,6 +29,7 @@ pub(crate) struct ClientInner { /// The runtime handle used for sending messages to the engine. handle: Handle, model_name: String, + scheduler_stats_recorder: SchedulerStatsRecorder, request_reg: Mutex, utility_reg: Mutex, health_error: ArcSwapOption, @@ -43,10 +44,13 @@ impl ClientInner { model_name: String, engines: &[ConnectedEngine], ) -> Self { + let scheduler_stats_recorder = + SchedulerStatsRecorder::new(&METRICS.scheduler, &model_name, engines); Self { input_send, handle, model_name, + scheduler_stats_recorder, request_reg: Mutex::new(RequestRegistry::new(engines)), utility_reg: Mutex::new(UtilityRegistry::default()), health_error: ArcSwapOption::empty(), @@ -389,12 +393,7 @@ pub(crate) async fn run_output_dispatcher_loop( "dropping scheduler stats for unknown engine" ); } - record_scheduler_stats( - &METRICS.scheduler, - inner.model_name(), - batch.engine_index, - scheduler_stats, - ); + inner.scheduler_stats_recorder.record(batch.engine_index, scheduler_stats); } // The engine's scheduler stats never carry adapter names; diff --git a/rust/src/engine-core-client/src/client/stream.rs b/rust/src/engine-core-client/src/client/stream.rs index 56c6a7cb663..b0ea180795a 100644 --- a/rust/src/engine-core-client/src/client/stream.rs +++ b/rust/src/engine-core-client/src/client/stream.rs @@ -45,6 +45,7 @@ impl Deref for EngineCoreStreamOutput { /// `finish_reason` is non-`None`. pub struct EngineCoreOutputStream { request_id: String, + engine_index: u32, abort_tx: mpsc::UnboundedSender, state: State, rx: OutputReceiver, @@ -53,11 +54,13 @@ pub struct EngineCoreOutputStream { impl EngineCoreOutputStream { pub(crate) fn new( request_id: String, + engine_index: u32, abort_tx: mpsc::UnboundedSender, rx: OutputReceiver, ) -> Self { Self { request_id, + engine_index, abort_tx, state: State::Running, rx, @@ -68,6 +71,11 @@ impl EngineCoreOutputStream { pub fn request_id(&self) -> &str { &self.request_id } + + /// Return the index of the engine that owns this request. + pub fn engine_index(&self) -> u32 { + self.engine_index + } } impl Stream for EngineCoreOutputStream { diff --git a/rust/src/engine-core-client/src/metrics.rs b/rust/src/engine-core-client/src/metrics.rs index e129fdd6a27..05744db42df 100644 --- a/rust/src/engine-core-client/src/metrics.rs +++ b/rust/src/engine-core-client/src/metrics.rs @@ -1,90 +1,190 @@ +use std::collections::BTreeMap; use std::collections::BTreeSet; use std::time::{SystemTime, UNIX_EPOCH}; use vllm_metrics::{ - EngineLabels, EnginePositionLabels, LoraAdapterNames, LoraInfoLabels, SchedulerMetrics, + EngineLabels, EnginePositionLabels, F64Gauge, Family, HistogramMetric, LoraAdapterNames, + LoraInfoLabels, SchedulerLogStatsAccumulator, SchedulerMetrics, U64Counter, U64Gauge, WaitingReasonLabels, }; use crate::protocol::stats::SchedulerStats; +use crate::transport::ConnectedEngine; const WAITING_REASON_CAPACITY: &str = "capacity"; const WAITING_REASON_DEFERRED: &str = "deferred"; -/// Record the scheduler-stats-backed metrics for one engine at one point in -/// time. -pub(crate) fn record_scheduler_stats( - metrics: &SchedulerMetrics, - model_name: impl Into, - engine: u32, - stats: &SchedulerStats, -) { - let model_name = model_name.into(); - let labels = EngineLabels { - model_name: model_name.clone(), - engine, - }; +/// Cached scheduler-stats metric handles for all engines connected to one +/// frontend client. +pub(crate) struct SchedulerStatsRecorder { + engines: BTreeMap, +} + +/// Per-engine cached metric handles used while recording `SchedulerStats`. +struct SchedulerStatsHandles { + // Base labels reused for dynamic child labels. + labels: EngineLabels, // Scheduler state gauges. - metrics.scheduler_running.get_or_create(&labels).set(stats.num_running_reqs); - metrics - .scheduler_waiting - .get_or_create(&labels) - .set(stats.num_waiting_reqs + stats.num_skipped_waiting_reqs); - metrics - .scheduler_waiting_by_reason - .get_or_create(&WaitingReasonLabels { - model_name: model_name.clone(), - engine, - reason: WAITING_REASON_CAPACITY, - }) - .set(stats.num_waiting_reqs); - metrics - .scheduler_waiting_by_reason - .get_or_create(&WaitingReasonLabels { - model_name: model_name.clone(), - engine, - reason: WAITING_REASON_DEFERRED, - }) - .set(stats.num_skipped_waiting_reqs); - metrics.kv_cache_usage.get_or_create(&labels).set(stats.kv_cache_usage); + scheduler_running: U64Gauge, + scheduler_waiting: U64Gauge, + scheduler_waiting_capacity: U64Gauge, + scheduler_waiting_deferred: U64Gauge, + kv_cache_usage: F64Gauge, // Prefix-cache counters, including the connector-backed external cache path. - metrics - .prefix_cache_queries - .get_or_create(&labels) - .inc_by(stats.prefix_cache_stats.base.queries); - metrics - .prefix_cache_hits - .get_or_create(&labels) - .inc_by(stats.prefix_cache_stats.base.hits); + prefix_cache_queries: U64Counter, + prefix_cache_hits: U64Counter, + external_prefix_cache_queries: U64Counter, + external_prefix_cache_hits: U64Counter, + + // Speculative decoding counters. + spec_decode_num_drafts: U64Counter, + spec_decode_num_draft_tokens: U64Counter, + spec_decode_num_accepted_tokens: U64Counter, + spec_decode_num_accepted_tokens_per_pos: Family, + + // Per-engine performance / MFU counters. + estimated_flops_per_gpu: U64Counter, + estimated_read_bytes_per_gpu: U64Counter, + estimated_write_bytes_per_gpu: U64Counter, + + // Sampled KV-cache residency histograms. + kv_block_lifetime_seconds: HistogramMetric, + kv_block_idle_before_evict_seconds: HistogramMetric, + kv_block_reuse_gap_seconds: HistogramMetric, + + // Non-Prometheus interval accumulator for periodic text-log helpers. + log_stats: SchedulerLogStatsAccumulator, +} + +impl SchedulerStatsRecorder { + /// Resolve the fixed-label metric handles for the connected engines. + pub(crate) fn new( + metrics: &SchedulerMetrics, + model_name: &str, + engines: &[ConnectedEngine], + ) -> Self { + let engines = engines + .iter() + .filter_map(|engine| { + let engine = engine.engine_id.engine_index()?; + Some(( + engine, + resolve_scheduler_stats_handles(metrics, model_name, engine), + )) + }) + .collect(); + + Self { engines } + } + + /// Record one scheduler-stats payload for the given engine index. + pub(crate) fn record(&self, engine_index: u32, stats: &SchedulerStats) { + if let Some(handles) = self.engines.get(&engine_index) { + record_scheduler_stats_with_handles(handles, stats); + } + } +} + +/// Resolve all fixed-label scheduler metrics for one engine. +fn resolve_scheduler_stats_handles( + metrics: &SchedulerMetrics, + model_name: &str, + engine: u32, +) -> SchedulerStatsHandles { + let labels = EngineLabels { + model_name: model_name.to_string(), + engine, + }; + let capacity = WaitingReasonLabels { + model_name: model_name.to_string(), + engine, + reason: WAITING_REASON_CAPACITY, + }; + let deferred = WaitingReasonLabels { + model_name: model_name.to_string(), + engine, + reason: WAITING_REASON_DEFERRED, + }; + + SchedulerStatsHandles { + scheduler_running: metrics.scheduler_running.get_or_create_owned(&labels), + scheduler_waiting: metrics.scheduler_waiting.get_or_create_owned(&labels), + scheduler_waiting_capacity: metrics + .scheduler_waiting_by_reason + .get_or_create_owned(&capacity), + scheduler_waiting_deferred: metrics + .scheduler_waiting_by_reason + .get_or_create_owned(&deferred), + kv_cache_usage: metrics.kv_cache_usage.get_or_create_owned(&labels), + prefix_cache_queries: metrics.prefix_cache_queries.get_or_create_owned(&labels), + prefix_cache_hits: metrics.prefix_cache_hits.get_or_create_owned(&labels), + external_prefix_cache_queries: metrics + .external_prefix_cache_queries + .get_or_create_owned(&labels), + external_prefix_cache_hits: metrics.external_prefix_cache_hits.get_or_create_owned(&labels), + spec_decode_num_drafts: metrics.spec_decode_num_drafts.get_or_create_owned(&labels), + spec_decode_num_draft_tokens: metrics + .spec_decode_num_draft_tokens + .get_or_create_owned(&labels), + spec_decode_num_accepted_tokens: metrics + .spec_decode_num_accepted_tokens + .get_or_create_owned(&labels), + spec_decode_num_accepted_tokens_per_pos: metrics + .spec_decode_num_accepted_tokens_per_pos + .clone(), + log_stats: metrics.log_stats.get_or_create_owned(&labels), + estimated_flops_per_gpu: metrics.estimated_flops_per_gpu.get_or_create_owned(&labels), + estimated_read_bytes_per_gpu: metrics + .estimated_read_bytes_per_gpu + .get_or_create_owned(&labels), + estimated_write_bytes_per_gpu: metrics + .estimated_write_bytes_per_gpu + .get_or_create_owned(&labels), + kv_block_lifetime_seconds: metrics.kv_block_lifetime_seconds.get_or_create_owned(&labels), + kv_block_idle_before_evict_seconds: metrics + .kv_block_idle_before_evict_seconds + .get_or_create_owned(&labels), + kv_block_reuse_gap_seconds: metrics.kv_block_reuse_gap_seconds.get_or_create_owned(&labels), + labels, + } +} + +/// Record scheduler-stats values through pre-resolved metric handles. +fn record_scheduler_stats_with_handles(handles: &SchedulerStatsHandles, stats: &SchedulerStats) { + // Scheduler state gauges. + handles.scheduler_running.set(stats.num_running_reqs); + handles + .scheduler_waiting + .set(stats.num_waiting_reqs + stats.num_skipped_waiting_reqs); + handles.scheduler_waiting_capacity.set(stats.num_waiting_reqs); + handles.scheduler_waiting_deferred.set(stats.num_skipped_waiting_reqs); + handles.kv_cache_usage.set(stats.kv_cache_usage); + + // Prefix-cache counters, including the connector-backed external cache path. + handles.prefix_cache_queries.inc_by(stats.prefix_cache_stats.base.queries); + handles.prefix_cache_hits.inc_by(stats.prefix_cache_stats.base.hits); if let Some(connector_prefix_cache_stats) = &stats.connector_prefix_cache_stats { - metrics + handles .external_prefix_cache_queries - .get_or_create(&labels) .inc_by(connector_prefix_cache_stats.base.queries); - metrics + handles .external_prefix_cache_hits - .get_or_create(&labels) .inc_by(connector_prefix_cache_stats.base.hits); } // Speculative decoding counters. if let Some(spec_decoding_stats) = &stats.spec_decoding_stats { - metrics - .spec_decode_num_drafts - .get_or_create(&labels) - .inc_by(spec_decoding_stats.num_drafts); - metrics + handles.spec_decode_num_drafts.inc_by(spec_decoding_stats.num_drafts); + handles .spec_decode_num_draft_tokens - .get_or_create(&labels) .inc_by(spec_decoding_stats.num_draft_tokens); - metrics + handles .spec_decode_num_accepted_tokens - .get_or_create(&labels) .inc_by(spec_decoding_stats.num_accepted_tokens); - metrics.log_stats.get_or_create(&labels).observe_spec_decode( + handles.log_stats.observe_spec_decode( spec_decoding_stats.num_drafts, &spec_decoding_stats.num_accepted_tokens_per_pos, ); @@ -92,11 +192,11 @@ pub(crate) fn record_scheduler_stats( for (position, accepted_tokens) in spec_decoding_stats.num_accepted_tokens_per_pos.iter().copied().enumerate() { - metrics + handles .spec_decode_num_accepted_tokens_per_pos .get_or_create(&EnginePositionLabels { - model_name: model_name.clone(), - engine, + model_name: handles.labels.model_name.clone(), + engine: handles.labels.engine, position: position as u32, }) .inc_by(accepted_tokens); @@ -109,22 +209,13 @@ pub(crate) fn record_scheduler_stats( || perf_stats.num_read_bytes_per_gpu != 0 || perf_stats.num_write_bytes_per_gpu != 0) { - metrics - .estimated_flops_per_gpu - .get_or_create(&labels) - .inc_by(perf_stats.num_flops_per_gpu); - metrics - .estimated_read_bytes_per_gpu - .get_or_create(&labels) - .inc_by(perf_stats.num_read_bytes_per_gpu); - metrics - .estimated_write_bytes_per_gpu - .get_or_create(&labels) - .inc_by(perf_stats.num_write_bytes_per_gpu); + handles.estimated_flops_per_gpu.inc_by(perf_stats.num_flops_per_gpu); + handles.estimated_read_bytes_per_gpu.inc_by(perf_stats.num_read_bytes_per_gpu); + handles.estimated_write_bytes_per_gpu.inc_by(perf_stats.num_write_bytes_per_gpu); } if let Some(cudagraph_stats) = &stats.cudagraph_stats { - metrics.log_stats.get_or_create(&labels).observe_cudagraph( + handles.log_stats.observe_cudagraph( cudagraph_stats.num_unpadded_tokens, cudagraph_stats.num_padded_tokens, cudagraph_stats.num_paddings, @@ -134,16 +225,11 @@ pub(crate) fn record_scheduler_stats( // Sampled KV-cache residency histograms. if !stats.kv_cache_eviction_events.is_empty() { - let kv_block_lifetime_seconds = metrics.kv_block_lifetime_seconds.get_or_create(&labels); - let kv_block_idle_before_evict_seconds = - metrics.kv_block_idle_before_evict_seconds.get_or_create(&labels); - let kv_block_reuse_gap_seconds = metrics.kv_block_reuse_gap_seconds.get_or_create(&labels); - for event in &stats.kv_cache_eviction_events { - kv_block_lifetime_seconds.observe(event.lifetime_seconds); - kv_block_idle_before_evict_seconds.observe(event.idle_seconds); + handles.kv_block_lifetime_seconds.observe(event.lifetime_seconds); + handles.kv_block_idle_before_evict_seconds.observe(event.idle_seconds); for reuse_gap_seconds in &event.reuse_gaps_seconds { - kv_block_reuse_gap_seconds.observe(*reuse_gap_seconds); + handles.kv_block_reuse_gap_seconds.observe(*reuse_gap_seconds); } } } diff --git a/rust/src/engine-core-client/src/protocol/tensor.rs b/rust/src/engine-core-client/src/protocol/tensor.rs index b6711215481..b80472129b7 100644 --- a/rust/src/engine-core-client/src/protocol/tensor.rs +++ b/rust/src/engine-core-client/src/protocol/tensor.rs @@ -11,6 +11,21 @@ use serde_tuple::{Deserialize_tuple, Serialize_tuple}; /// const CUSTOM_TYPE_RAW_VIEW: i8 = 3; +#[derive(Serialize)] +#[serde(rename = "_ExtStruct")] +struct MsgpackExtRef<'a>((i8, ByteSlice<'a>)); + +struct ByteSlice<'a>(&'a [u8]); + +impl Serialize for ByteSlice<'_> { + fn serialize(&self, serializer: S) -> std::result::Result + where + S: Serializer, + { + serializer.serialize_bytes(self.0) + } +} + #[easy_ext::ext(ShapeExt)] impl [usize] { /// Returned the total number of elements implied by this shape, or `None` @@ -184,7 +199,7 @@ impl Serialize for WireArrayData { match self { Self::AuxIndex(index) => serializer.serialize_u64(*index as u64), Self::RawView(bytes) => { - Value::Ext(CUSTOM_TYPE_RAW_VIEW, bytes.clone()).serialize(serializer) + MsgpackExtRef((CUSTOM_TYPE_RAW_VIEW, ByteSlice(bytes))).serialize(serializer) } } } @@ -194,6 +209,21 @@ impl Serialize for WireArrayData { mod tests { use super::*; + #[test] + fn raw_view_serializes_as_msgpack_ext() { + let bytes = vec![1, 2, 3, 4]; + let encoded = + rmp_serde::to_vec_named(&WireArrayData::RawView(bytes.clone())).expect("encode"); + let expected = rmp_serde::to_vec_named(&Value::Ext(CUSTOM_TYPE_RAW_VIEW, bytes.clone())) + .expect("encode expected"); + + assert_eq!(encoded, expected); + assert_eq!( + rmpv::decode::read_value(&mut std::io::Cursor::new(encoded)).expect("decode"), + Value::Ext(CUSTOM_TYPE_RAW_VIEW, bytes) + ); + } + #[test] fn constructors_build_raw_view_tensors() { let f32_tensor = WireNdArray::from_f32(vec![2], vec![1.0, 2.5]).unwrap(); diff --git a/rust/src/llm/src/lib.rs b/rust/src/llm/src/lib.rs index ce15a970c63..942bf55c288 100644 --- a/rust/src/llm/src/lib.rs +++ b/rust/src/llm/src/lib.rs @@ -14,6 +14,7 @@ pub use output::{ GenerateOutputStreamExt, GeneratePromptInfo, TokenUsage, }; pub use request::GenerateRequest; +pub use request_metrics::current_unix_timestamp_secs; pub use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs, TokenLogprob}; use crate::inflight::InflightRequests; @@ -88,14 +89,21 @@ impl Llm { // Record internal engine-core request ID in the current tracing span. Span::current().record("engine_request_id", &internal_request_id); + let arrival_time = prepared.engine_request.arrival_time; + let max_tokens_param = + (prepared.engine_request.sampling_params.as_ref()).map(|p| p.max_tokens); + let prompt_len = prepared.prompt_token_ids().len() as u32; + + let stream = self.client.call(prepared.engine_request).await?; + let request_metrics = RequestMetricsTracker::new( self.client.model_name().to_string(), - prepared.engine_request.arrival_time, - prepared.prompt_token_ids().len() as u32, - (prepared.engine_request.sampling_params.as_ref()).map(|p| p.max_tokens), + stream.engine_index(), + arrival_time, + prompt_len, + max_tokens_param, 1, ); - let stream = self.client.call(prepared.engine_request).await?; let guard = self.inflight.track(external_request_id, internal_request_id); Ok(GenerateOutputStream::new( diff --git a/rust/src/llm/src/output.rs b/rust/src/llm/src/output.rs index 609088ec6b2..d1d7e5f46e5 100644 --- a/rust/src/llm/src/output.rs +++ b/rust/src/llm/src/output.rs @@ -248,12 +248,7 @@ impl Stream for GenerateOutputStream { }; let received_at = current_unix_timestamp_secs(); - self.request_metrics.observe_output( - raw.engine_index, - raw.timestamp, - received_at, - &raw.output, - ); + self.request_metrics.observe_output(raw.timestamp, received_at, &raw.output); let raw = raw.output; diff --git a/rust/src/llm/src/request.rs b/rust/src/llm/src/request.rs index 159cf823de4..45e5bd1ca64 100644 --- a/rust/src/llm/src/request.rs +++ b/rust/src/llm/src/request.rs @@ -1,5 +1,4 @@ use std::collections::BTreeMap; -use std::time::{SystemTime, UNIX_EPOCH}; use uuid::Uuid; use vllm_engine_core_client::protocol::lora::LoraRequest; @@ -8,6 +7,7 @@ use vllm_engine_core_client::protocol::request::{EngineCoreRequest, ReasoningPar use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; use crate::error::{Error, Result}; +use crate::request_metrics::current_unix_timestamp_secs; /// Tokenized decoder-only generate request accepted by [`crate::Llm`]. /// @@ -30,8 +30,9 @@ pub struct GenerateRequest { pub mm_features: Option, /// Unix timestamp, in seconds, when this request arrived at the frontend. /// - /// When omitted, the Rust frontend fills it immediately before sending the - /// request to engine-core, matching Python's default arrival-time behavior. + /// Stamped at the frontend entry, before render and tokenization, to match + /// Python's renderer-entry arrival_time. When omitted, it is filled as a + /// fallback before the request is sent to engine-core. pub arrival_time: Option, /// Optional salt used to partition prefix-cache entries for this request. pub cache_salt: Option, @@ -122,13 +123,6 @@ impl PreparedGenerateRequest { } } -fn current_unix_timestamp_secs() -> f64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .expect("system clock is before unix epoch") - .as_secs_f64() -} - #[cfg(test)] mod tests { use std::collections::BTreeMap; diff --git a/rust/src/llm/src/request_metrics.rs b/rust/src/llm/src/request_metrics.rs index 38795a70928..4f1673db154 100644 --- a/rust/src/llm/src/request_metrics.rs +++ b/rust/src/llm/src/request_metrics.rs @@ -5,15 +5,12 @@ use vllm_engine_core_client::protocol::output::{ }; use vllm_engine_core_client::protocol::stats::PrefillStats; use vllm_metrics::{ - EngineLabels, FinishedReasonLabels, METRICS, PromptTokenSourceLabels, RequestMetrics, + EngineLabels, Family, FinishedReasonLabels, HistogramMetric, METRICS, PromptTokenSourceLabels, + U64Counter, }; use crate::FinishReason; -fn metrics() -> &'static RequestMetrics { - &METRICS.request -} - const PROMPT_TOKEN_SOURCE_LOCAL_COMPUTE: &str = "local_compute"; const PROMPT_TOKEN_SOURCE_LOCAL_CACHE_HIT: &str = "local_cache_hit"; const PROMPT_TOKEN_SOURCE_EXTERNAL_KV_TRANSFER: &str = "external_kv_transfer"; @@ -29,9 +26,11 @@ const PROMPT_TOKEN_SOURCE_EXTERNAL_KV_TRANSFER: &str = "external_kv_transfer"; /// /// Original Python update flow: /// -#[derive(Debug, Clone)] +#[derive(Clone)] pub(crate) struct RequestMetricsTracker { - model_name: String, + /// Cached request metric handles for this request's model and engine index. + handles: RequestMetricHandles, + arrival_time: f64, prompt_len: u32, max_tokens_param: Option, @@ -44,7 +43,38 @@ pub(crate) struct RequestMetricsTracker { first_token_latency: f64, num_generation_tokens: u32, latest_num_cached_tokens: u32, - last_seen_engine_index: u32, +} + +/// Cached request metric handles for one model and engine index. +#[derive(Clone)] +struct RequestMetricHandles { + labels: EngineLabels, + + // Request-derived counters. + num_preemptions: U64Counter, + prompt_tokens: U64Counter, + prompt_tokens_local_compute: U64Counter, + prompt_tokens_local_cache_hit: U64Counter, + prompt_tokens_external_kv_transfer: U64Counter, + prompt_tokens_cached: U64Counter, + generation_tokens: U64Counter, + + // Request lifecycle counters and histograms. + request_success: Family, + request_prompt_tokens: HistogramMetric, + request_generation_tokens: HistogramMetric, + request_max_num_generation_tokens: HistogramMetric, + request_params_max_tokens: HistogramMetric, + request_params_n: HistogramMetric, + request_prefill_kv_computed_tokens: HistogramMetric, + time_to_first_token_seconds: HistogramMetric, + inter_token_latency_seconds: HistogramMetric, + e2e_request_latency_seconds: HistogramMetric, + request_queue_time_seconds: HistogramMetric, + request_prefill_time_seconds: HistogramMetric, + request_decode_time_seconds: HistogramMetric, + request_inference_time_seconds: HistogramMetric, + request_time_per_output_token_seconds: HistogramMetric, } impl RequestMetricsTracker { @@ -52,13 +82,14 @@ impl RequestMetricsTracker { /// context. pub(crate) fn new( model_name: String, + engine_index: u32, arrival_time: f64, prompt_len: u32, max_tokens_param: Option, n_param: u32, ) -> Self { Self { - model_name, + handles: resolve_request_metric_handles(&model_name, engine_index), arrival_time, prompt_len, max_tokens_param, @@ -71,7 +102,6 @@ impl RequestMetricsTracker { first_token_latency: 0.0, num_generation_tokens: 0, latest_num_cached_tokens: 0, - last_seen_engine_index: 0, } } @@ -81,23 +111,18 @@ impl RequestMetricsTracker { /// pub(crate) fn observe_output( &mut self, - engine_index: u32, batch_timestamp: f64, received_at: f64, output: &EngineCoreOutput, ) { - self.last_seen_engine_index = engine_index; if let Some(prefill_stats) = &output.prefill_stats { self.latest_num_cached_tokens = prefill_stats.num_cached_tokens; } self.num_generation_tokens += output.new_token_ids.len() as u32; - metrics() - .generation_tokens - .get_or_create(&engine_labels(&self.model_name, engine_index)) - .inc_by(output.new_token_ids.len() as u64); + self.handles.generation_tokens.inc_by(output.new_token_ids.len() as u64); if let Some(events) = &output.events { - self.observe_events(engine_index, events); + self.observe_events(events); } // Only outputs that actually carry tokens drive token-timing metrics. @@ -107,22 +132,16 @@ impl RequestMetricsTracker { if !output.new_token_ids.is_empty() { if self.is_prefilling { if let Some(prefill_stats) = &output.prefill_stats { - record_prompt_tokens(&self.model_name, engine_index, prefill_stats); + self.record_prompt_tokens(prefill_stats); } self.first_token_latency = received_at - self.arrival_time; - observe_time_to_first_token_seconds( - &self.model_name, - engine_index, - self.first_token_latency, - ); + self.handles.time_to_first_token_seconds.observe(self.first_token_latency); self.first_token_ts = batch_timestamp; self.is_prefilling = false; } else if self.last_token_ts > 0.0 { - observe_inter_token_latency_seconds( - &self.model_name, - engine_index, - batch_timestamp - self.last_token_ts, - ); + self.handles + .inter_token_latency_seconds + .observe(batch_timestamp - self.last_token_ts); } self.last_token_ts = batch_timestamp; @@ -135,7 +154,6 @@ impl RequestMetricsTracker { /// Original Python finished-request stats: /// pub(crate) fn record_finished(&self, received_at: f64, finish_reason: FinishReason) { - let labels = engine_labels(&self.model_name, self.last_seen_engine_index); let prefill_kv_computed_tokens = self.prompt_len.saturating_sub(self.latest_num_cached_tokens); let e2e_latency_seconds = received_at - self.arrival_time; @@ -150,57 +168,47 @@ impl RequestMetricsTracker { 0.0 }; - record_request_success(&self.model_name, self.last_seen_engine_index, finish_reason); - metrics() - .request_prompt_tokens - .get_or_create(&labels) - .observe(self.prompt_len as f64); - metrics() + self.record_request_success(finish_reason); + + self.handles.request_prompt_tokens.observe(self.prompt_len as f64); + self.handles .request_generation_tokens - .get_or_create(&labels) .observe(self.num_generation_tokens as f64); - metrics() + self.handles .request_max_num_generation_tokens - .get_or_create(&labels) .observe(self.num_generation_tokens as f64); if let Some(max_tokens_param) = self.max_tokens_param { - metrics() - .request_params_max_tokens - .get_or_create(&labels) - .observe(max_tokens_param as f64); + self.handles.request_params_max_tokens.observe(max_tokens_param as f64); } - metrics().request_params_n.get_or_create(&labels).observe(self.n_param as f64); - metrics() + self.handles.request_params_n.observe(self.n_param as f64); + self.handles .request_prefill_kv_computed_tokens - .get_or_create(&labels) .observe(prefill_kv_computed_tokens as f64); - metrics() - .e2e_request_latency_seconds - .get_or_create(&labels) - .observe(e2e_latency_seconds); - metrics() - .request_queue_time_seconds - .get_or_create(&labels) - .observe(queue_time_seconds); - metrics() - .request_prefill_time_seconds - .get_or_create(&labels) - .observe(prefill_time_seconds); - metrics() - .request_decode_time_seconds - .get_or_create(&labels) - .observe(decode_time_seconds); - metrics() - .request_inference_time_seconds - .get_or_create(&labels) - .observe(inference_time_seconds); - metrics() + self.handles.e2e_request_latency_seconds.observe(e2e_latency_seconds); + self.handles.request_queue_time_seconds.observe(queue_time_seconds); + self.handles.request_prefill_time_seconds.observe(prefill_time_seconds); + self.handles.request_decode_time_seconds.observe(decode_time_seconds); + self.handles.request_inference_time_seconds.observe(inference_time_seconds); + self.handles .request_time_per_output_token_seconds - .get_or_create(&labels) .observe(time_per_output_token_seconds); } - fn observe_events(&mut self, engine_index: u32, events: &[EngineCoreEvent]) { + /// Record prompt token counters through cached metric handles. + fn record_prompt_tokens(&self, prefill_stats: &PrefillStats) { + let computed = prefill_stats.num_computed_tokens as u64; + let local_cache_hit = prefill_stats.num_local_cached_tokens as u64; + let external_kv_transfer = prefill_stats.num_external_cached_tokens as u64; + + self.handles.prompt_tokens.inc_by(prefill_stats.num_prompt_tokens as u64); + self.handles.prompt_tokens_local_compute.inc_by(computed); + self.handles.prompt_tokens_local_cache_hit.inc_by(local_cache_hit); + self.handles.prompt_tokens_external_kv_transfer.inc_by(external_kv_transfer); + self.handles.prompt_tokens_cached.inc_by(prefill_stats.num_cached_tokens as u64); + } + + /// Record request event counters through cached metric handles. + fn observe_events(&mut self, events: &[EngineCoreEvent]) { for event in events { match event.r#type { EngineCoreEventType::Queued => { @@ -212,46 +220,86 @@ impl RequestMetricsTracker { } } EngineCoreEventType::Preempted => { - metrics() - .num_preemptions - .get_or_create(&engine_labels(&self.model_name, engine_index)) - .inc(); + self.handles.num_preemptions.inc(); } } } } -} -fn engine_labels(model_name: &str, engine: u32) -> EngineLabels { - EngineLabels { - model_name: model_name.to_string(), - engine, + /// Increment the request-success counter for the terminal finish reason. + fn record_request_success(&self, finish_reason: FinishReason) { + self.handles + .request_success + .get_or_create(&FinishedReasonLabels { + model_name: self.handles.labels.model_name.clone(), + engine: self.handles.labels.engine, + finished_reason: finish_reason.as_str(), + }) + .inc(); } } -fn observe_time_to_first_token_seconds(model_name: &str, engine: u32, seconds: f64) { - metrics() - .time_to_first_token_seconds - .get_or_create(&engine_labels(model_name, engine)) - .observe(seconds); -} +/// Resolve fixed request metric handles for one model and engine index. +fn resolve_request_metric_handles(model_name: &str, engine: u32) -> RequestMetricHandles { + let metrics = &METRICS.request; + let labels = EngineLabels { + model_name: model_name.to_string(), + engine, + }; -fn observe_inter_token_latency_seconds(model_name: &str, engine: u32, seconds: f64) { - metrics() - .inter_token_latency_seconds - .get_or_create(&engine_labels(model_name, engine)) - .observe(seconds); -} - -fn record_request_success(model_name: &str, engine: u32, finish_reason: FinishReason) { - metrics() - .request_success - .get_or_create(&FinishedReasonLabels { - model_name: model_name.to_string(), - engine, - finished_reason: finish_reason.as_str(), - }) - .inc(); + RequestMetricHandles { + num_preemptions: metrics.num_preemptions.get_or_create_owned(&labels), + prompt_tokens: metrics.prompt_tokens.get_or_create_owned(&labels), + prompt_tokens_local_compute: metrics.prompt_tokens_by_source.get_or_create_owned( + &prompt_token_source_labels(model_name, engine, PROMPT_TOKEN_SOURCE_LOCAL_COMPUTE), + ), + prompt_tokens_local_cache_hit: metrics.prompt_tokens_by_source.get_or_create_owned( + &prompt_token_source_labels(model_name, engine, PROMPT_TOKEN_SOURCE_LOCAL_CACHE_HIT), + ), + prompt_tokens_external_kv_transfer: metrics.prompt_tokens_by_source.get_or_create_owned( + &prompt_token_source_labels( + model_name, + engine, + PROMPT_TOKEN_SOURCE_EXTERNAL_KV_TRANSFER, + ), + ), + prompt_tokens_cached: metrics.prompt_tokens_cached.get_or_create_owned(&labels), + generation_tokens: metrics.generation_tokens.get_or_create_owned(&labels), + request_success: metrics.request_success.clone(), + request_prompt_tokens: metrics.request_prompt_tokens.get_or_create_owned(&labels), + request_generation_tokens: metrics.request_generation_tokens.get_or_create_owned(&labels), + request_max_num_generation_tokens: metrics + .request_max_num_generation_tokens + .get_or_create_owned(&labels), + request_params_max_tokens: metrics.request_params_max_tokens.get_or_create_owned(&labels), + request_params_n: metrics.request_params_n.get_or_create_owned(&labels), + request_prefill_kv_computed_tokens: metrics + .request_prefill_kv_computed_tokens + .get_or_create_owned(&labels), + time_to_first_token_seconds: metrics + .time_to_first_token_seconds + .get_or_create_owned(&labels), + inter_token_latency_seconds: metrics + .inter_token_latency_seconds + .get_or_create_owned(&labels), + e2e_request_latency_seconds: metrics + .e2e_request_latency_seconds + .get_or_create_owned(&labels), + request_queue_time_seconds: metrics.request_queue_time_seconds.get_or_create_owned(&labels), + request_prefill_time_seconds: metrics + .request_prefill_time_seconds + .get_or_create_owned(&labels), + request_decode_time_seconds: metrics + .request_decode_time_seconds + .get_or_create_owned(&labels), + request_inference_time_seconds: metrics + .request_inference_time_seconds + .get_or_create_owned(&labels), + request_time_per_output_token_seconds: metrics + .request_time_per_output_token_seconds + .get_or_create_owned(&labels), + labels, + } } fn prompt_token_source_labels( @@ -266,45 +314,6 @@ fn prompt_token_source_labels( } } -fn record_prompt_tokens(model_name: &str, engine: u32, prefill_stats: &PrefillStats) { - let computed = prefill_stats.num_computed_tokens as u64; - let local_cache_hit = prefill_stats.num_local_cached_tokens as u64; - let external_kv_transfer = prefill_stats.num_external_cached_tokens as u64; - - metrics() - .prompt_tokens - .get_or_create(&engine_labels(model_name, engine)) - .inc_by(prefill_stats.num_prompt_tokens as u64); - metrics() - .prompt_tokens_by_source - .get_or_create(&prompt_token_source_labels( - model_name, - engine, - PROMPT_TOKEN_SOURCE_LOCAL_COMPUTE, - )) - .inc_by(computed); - metrics() - .prompt_tokens_by_source - .get_or_create(&prompt_token_source_labels( - model_name, - engine, - PROMPT_TOKEN_SOURCE_LOCAL_CACHE_HIT, - )) - .inc_by(local_cache_hit); - metrics() - .prompt_tokens_by_source - .get_or_create(&prompt_token_source_labels( - model_name, - engine, - PROMPT_TOKEN_SOURCE_EXTERNAL_KV_TRANSFER, - )) - .inc_by(external_kv_transfer); - metrics() - .prompt_tokens_cached - .get_or_create(&engine_labels(model_name, engine)) - .inc_by(prefill_stats.num_cached_tokens as u64); -} - fn diff_or_zero(end: f64, start: f64) -> f64 { if end > 0.0 && start > 0.0 && end >= start { end - start @@ -321,7 +330,7 @@ fn diff_or_zero(end: f64, start: f64) -> f64 { /// /// Original Python request timestamp source: /// -pub(crate) fn current_unix_timestamp_secs() -> f64 { +pub fn current_unix_timestamp_secs() -> f64 { SystemTime::now() .duration_since(UNIX_EPOCH) .expect("system clock is before unix epoch") @@ -337,10 +346,10 @@ mod tests { #[test] fn tracker_updates_timing_state_across_prefill_decode_and_finish() { - let mut tracker = RequestMetricsTracker::new("model".to_string(), 100.0, 64, Some(128), 1); + let mut tracker = + RequestMetricsTracker::new("model".to_string(), 2, 100.0, 64, Some(128), 1); tracker.observe_output( - 2, 10.0, 100.2, &vllm_engine_core_client::protocol::output::EngineCoreOutput { @@ -368,7 +377,6 @@ mod tests { }, ); tracker.observe_output( - 2, 11.5, 100.4, &vllm_engine_core_client::protocol::output::EngineCoreOutput { @@ -384,7 +392,7 @@ mod tests { ); assert!(!tracker.is_prefilling); - assert_eq!(tracker.last_seen_engine_index, 2); + assert_eq!(tracker.handles.labels.engine, 2); assert_eq!(tracker.num_generation_tokens, 3); assert_eq!(tracker.queued_ts, 8.0); assert_eq!(tracker.scheduled_ts, 9.0); diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs index 9911841eade..8581b1ac08f 100644 --- a/rust/src/llm/tests/generate.rs +++ b/rust/src/llm/tests/generate.rs @@ -17,7 +17,7 @@ use vllm_engine_core_client::protocol::request::EngineCoreRequest; use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; use vllm_engine_core_client::protocol::stats::PrefillStats; use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task}; -use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig}; +use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId}; use vllm_llm::{ Error, FinishReason, GenerateOutputStreamExt as _, GeneratePromptInfo, GenerateRequest, Llm, }; @@ -699,7 +699,7 @@ async fn abort_by_external_id_aborts_all_internal_requests() { async fn generate_records_request_metrics_in_prometheus_output() { let ipc = IpcNamespace::new().unwrap(); let handshake_address = ipc.handshake_endpoint(); - let engine_id = b"engine-metrics".to_vec(); + let engine_id = EngineId::from_engine_index(4); let model_name = request_metrics_model_name("metrics-model"); let (shutdown_tx, engine_task) = spawn_mock_engine_task( @@ -832,7 +832,7 @@ async fn generate_records_request_metrics_in_prometheus_output() { async fn dropping_stream_records_abort_terminal_request_metrics() { let ipc = IpcNamespace::new().unwrap(); let handshake_address = ipc.handshake_endpoint(); - let engine_id = b"engine-metrics-drop".to_vec(); + let engine_id = EngineId::from_engine_index(5); let model_name = request_metrics_model_name("metrics-drop-model"); let (shutdown_tx, engine_task) = spawn_mock_engine_task( diff --git a/rust/src/metrics/src/lib.rs b/rust/src/metrics/src/lib.rs index 8f0db53d3ff..ca650fbadae 100644 --- a/rust/src/metrics/src/lib.rs +++ b/rust/src/metrics/src/lib.rs @@ -4,7 +4,7 @@ use std::sync::atomic::AtomicU64; use prometheus_client::encoding::text::encode; use prometheus_client::metrics::counter::Counter; -use prometheus_client::metrics::family::Family; +pub use prometheus_client::metrics::family::Family; use prometheus_client::metrics::gauge::Gauge; use prometheus_client::metrics::histogram::Histogram; use prometheus_client::registry::Registry; @@ -23,6 +23,8 @@ pub use scheduler::*; pub type U64Counter = Counter; pub type U64Gauge = Gauge; pub type F64Gauge = Gauge; +/// Histogram metric handle cloned out of a Prometheus family. +pub type HistogramMetric = Histogram; pub(crate) type HistogramFamily = Family Histogram>; /// Shared Prometheus registry for frontend metrics. diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index 9de89727205..4327221221d 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -94,6 +94,7 @@ pub fn to_text_request( data_parallel_rank: None, reasoning_parser_kwargs: None, lora_request: None, + arrival_time: None, }) } diff --git a/rust/src/server/src/routes/inference/generate/convert.rs b/rust/src/server/src/routes/inference/generate/convert.rs index 7a6dcdfc45c..965155b5825 100644 --- a/rust/src/server/src/routes/inference/generate/convert.rs +++ b/rust/src/server/src/routes/inference/generate/convert.rs @@ -70,6 +70,7 @@ pub(super) fn prepare_generate_request( data_parallel_rank: ctx.data_parallel_rank, reasoning_parser_kwargs: None, lora_request: lora_resolution.lora_request.clone(), + arrival_time: None, }; Ok(PreparedRequest { diff --git a/rust/src/server/src/routes/openai/completions/convert.rs b/rust/src/server/src/routes/openai/completions/convert.rs index fa80e16e0f9..1355481b49b 100644 --- a/rust/src/server/src/routes/openai/completions/convert.rs +++ b/rust/src/server/src/routes/openai/completions/convert.rs @@ -144,6 +144,7 @@ pub(super) fn prepare_completion_request( data_parallel_rank: ctx.data_parallel_rank, reasoning_parser_kwargs: None, lora_request: lora_resolution.lora_request.clone(), + arrival_time: None, }; Ok(PreparedRequest { diff --git a/rust/src/text/src/lib.rs b/rust/src/text/src/lib.rs index 2987ef93e57..130eaec7f42 100644 --- a/rust/src/text/src/lib.rs +++ b/rust/src/text/src/lib.rs @@ -132,6 +132,10 @@ impl TextLlm { ) -> Result<(TextRequest, GenerateOutputStream)> { request.validate()?; + if request.arrival_time.is_none() { + request.arrival_time = Some(vllm_llm::current_unix_timestamp_secs()); + } + let tokenizer = self.backend.tokenizer(); let prompt_token_ids = match take(&mut request.prompt) { Prompt::Text(text) => tokenizer.encode(&text, request.add_special_tokens)?, diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index eabece37f09..164d2a3db06 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -18,7 +18,7 @@ use crate::request::{SamplingParams, TextRequest}; /// One text request after it has been lowered into the raw generate boundary. #[derive(Debug)] pub struct PreparedTextRequest { - /// The original high-level request, preserved for response-side metadata + /// The high-level request fields still needed for response-side metadata /// and decoding options. pub text_request: TextRequest, /// The southbound request ready to be sent to `vllm-llm`. @@ -28,7 +28,7 @@ pub struct PreparedTextRequest { /// Convert a high-level [`TextRequest`] into one lower-level /// [`GenerateRequest`] ready for the `llm` crate. pub fn lower_text_request( - request: TextRequest, + mut request: TextRequest, prompt_token_ids: Vec, sampling_hints: SamplingHints, sampling_limits: SamplingLimits, @@ -40,7 +40,10 @@ pub fn lower_text_request( let generate_request = GenerateRequest { request_id: request.request_id.clone(), prompt_token_ids, - mm_features: request.mm_features.clone(), + // Align with Python's response path: decoded output state does not retain + // `mm_features`; move them to the engine request to avoid cloning large + // multimodal tensor payloads. + mm_features: request.mm_features.take(), sampling_params: lower_sampling_params( request.sampling_params.clone(), sampling_hints, @@ -53,7 +56,7 @@ pub fn lower_text_request( data_parallel_rank: request.data_parallel_rank, reasoning_parser_kwargs: request.reasoning_parser_kwargs.clone(), lora_request: request.lora_request.clone(), - arrival_time: None, + arrival_time: request.arrival_time, trace_headers: None, }; @@ -307,6 +310,7 @@ mod tests { use std::collections::{BTreeSet, HashMap}; use serial_test::file_serial; + use vllm_engine_core_client::protocol::multimodal::{MmFeatureSpec, PlaceholderRange}; use vllm_tokenizer::test_utils::TestTokenizer; use super::*; @@ -574,6 +578,35 @@ mod tests { .assert_debug_eq(¶ms); } + #[test] + fn lower_text_request_moves_multimodal_features_to_generate_request() { + let features = vec![MmFeatureSpec { + data: None, + modality: "image".to_string(), + identifier: "image-1".to_string(), + mm_position: PlaceholderRange { + offset: 2, + length: 4, + is_embed: None, + }, + mm_hash: Some("hash-1".to_string()), + }]; + let mut request = sample_request(); + request.mm_features = Some(features.clone()); + + let prepared = lower_text_request( + request, + vec![1, 2, 3], + sample_sampling_hints(), + sample_sampling_limits(), + &stub_tokenizer(), + ) + .unwrap(); + + assert_eq!(prepared.generate_request.mm_features, Some(features)); + assert_eq!(prepared.text_request.mm_features, None); + } + #[test] fn lower_text_request_uses_union_vocab_for_prompt_token_ids() { lower_text_request( @@ -1110,6 +1143,44 @@ mod tests { assert_eq!(prepared.generate_request.request_id, "text-1"); } + #[test] + fn lower_text_request_passes_arrival_time_through() { + let request = TextRequest { + arrival_time: Some(42.5), + ..sample_request() + }; + + let prepared = lower_text_request( + request, + vec![1, 2, 3], + sample_sampling_hints(), + sample_sampling_limits(), + &stub_tokenizer(), + ) + .unwrap(); + + assert_eq!(prepared.generate_request.arrival_time, Some(42.5)); + } + + #[test] + fn lower_text_request_leaves_arrival_time_unset_when_absent() { + let request = TextRequest { + arrival_time: None, + ..sample_request() + }; + + let prepared = lower_text_request( + request, + vec![1, 2, 3], + sample_sampling_hints(), + sample_sampling_limits(), + &stub_tokenizer(), + ) + .unwrap(); + + assert_eq!(prepared.generate_request.arrival_time, None); + } + #[test] fn resolve_max_tokens_user_smaller_than_model_limit() { let result = resolve_max_tokens(Some(50), None, 200, 100); diff --git a/rust/src/text/src/request.rs b/rust/src/text/src/request.rs index c64e9ca05e6..09522868872 100644 --- a/rust/src/text/src/request.rs +++ b/rust/src/text/src/request.rs @@ -187,6 +187,12 @@ pub struct TextRequest { /// LoRA adapter selected for this request. #[serde(default)] pub lora_request: Option, + /// Wall-clock unix timestamp (seconds) when this request arrived at the + /// frontend, stamped before render/tokenize to match Python's + /// renderer-entry arrival_time. When unset, it is stamped before + /// tokenization. + #[serde(default)] + pub arrival_time: Option, } impl TextRequest { @@ -205,6 +211,7 @@ impl TextRequest { data_parallel_rank: None, reasoning_parser_kwargs: None, lora_request: None, + arrival_time: None, } } diff --git a/tests/compile/fusions_e2e/conftest.py b/tests/compile/fusions_e2e/conftest.py index 9f34d25c46d..a4ed63ffe7b 100644 --- a/tests/compile/fusions_e2e/conftest.py +++ b/tests/compile/fusions_e2e/conftest.py @@ -79,6 +79,7 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn): ): monkeypatch.setenv("VLLM_USE_DEEP_GEMM", "1" if use_deepgemm else "0") monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1" if use_aiter else "0") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_CUSTOM_AR", "1" if use_aiter else "0") from vllm._aiter_ops import rocm_aiter_ops rocm_aiter_ops.refresh_env_variables() diff --git a/tests/compile/passes/distributed/test_fusion_all_reduce.py b/tests/compile/passes/distributed/test_fusion_all_reduce.py index 83ce458aafd..b86018a7555 100644 --- a/tests/compile/passes/distributed/test_fusion_all_reduce.py +++ b/tests/compile/passes/distributed/test_fusion_all_reduce.py @@ -13,6 +13,7 @@ from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant from vllm.compilation.passes.fusion.allreduce_rms_fusion import ( AllReduceFusionPass, RocmAiterAllReduceFusionPass, + _select_flashinfer_allreduce_use_oneshot, ) from vllm.compilation.passes.fx_utils import find_op_nodes from vllm.compilation.passes.utility.fix_functionalization import ( @@ -30,6 +31,9 @@ from vllm.config import ( set_current_vllm_config, ) from vllm.distributed import tensor_model_parallel_all_reduce +from vllm.distributed.device_communicators.aiter_custom_all_reduce import ( + AiterCustomAllreduce, +) from vllm.distributed.parallel_state import ( init_distributed_environment, initialize_model_parallel, @@ -45,6 +49,35 @@ from vllm.utils.torch_utils import set_random_seed DEVICE_TYPE = current_platform.device_type +@pytest.mark.parametrize( + ("workspace_backend", "device_capability", "world_size", "tensor_size", "expected"), + [ + ("mnnvl", 103, 8, 2 * 1024 * 1024, None), + ("trtllm", 103, 8, 2 * 1024 * 1024, True), + ("trtllm", 103, 8, 2 * 1024 * 1024 + 1, False), + ("trtllm", 100, 4, 4 * 1024 * 1024, True), + ("trtllm", 100, 4, 4 * 1024 * 1024 + 1, False), + ("trtllm", None, 8, 128 * 1024 * 1024, True), + ], +) +def test_select_flashinfer_allreduce_use_oneshot( + workspace_backend: str, + device_capability: int | None, + world_size: int, + tensor_size: int, + expected: bool | None, +): + assert ( + _select_flashinfer_allreduce_use_oneshot( + workspace_backend, + device_capability, + world_size, + tensor_size, + ) + is expected + ) + + class TestAllReduceRMSNormModel(torch.nn.Module): def __init__( self, @@ -504,8 +537,12 @@ def all_reduce_fusion_pass_on_test_model( "MASTER_ADDR": "localhost", "MASTER_PORT": "12345", "VLLM_FLASHINFER_ALLREDUCE_BACKEND": flashinfer_allreduce_backend, + "VLLM_ROCM_USE_AITER": str(int(use_aiter)), + "VLLM_ROCM_USE_AITER_CUSTOM_AR": str(int(use_aiter)), } ) + if use_aiter: + rocm_aiter_ops.refresh_env_variables() init_distributed_environment() @@ -616,7 +653,7 @@ def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace( m.setenv("VLLM_ROCM_USE_AITER", "1") rocm_aiter_ops.refresh_env_variables() - if not rocm_aiter_ops.has_fused_allreduce_rmsnorm_quant_per_group(): + if not AiterCustomAllreduce.build_supports_per_group_quant(): pytest.skip( "aiter build is missing 'fused_ar_rms_per_group_quant' (needs " "ROCm/aiter PR #2823); the new patterns aren't registered." @@ -671,6 +708,7 @@ def rocm_aiter_group_quant_fusion_pass_on_test_model( "MASTER_ADDR": "localhost", "MASTER_PORT": "12345", "VLLM_ROCM_USE_AITER": "1", + "VLLM_ROCM_USE_AITER_CUSTOM_AR": "1", } ) rocm_aiter_ops.refresh_env_variables() diff --git a/tests/config/test_bailing_mtp_config.py b/tests/config/test_bailing_mtp_config.py new file mode 100644 index 00000000000..8fae29959f2 --- /dev/null +++ b/tests/config/test_bailing_mtp_config.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from transformers import PretrainedConfig + +from vllm.config.speculative import MTPModelTypes, SpeculativeConfig +from vllm.transformers_utils.model_arch_config_convertor import ( + BailingHybridMTPModelArchConfigConvertor, +) + + +def _bailing_config() -> PretrainedConfig: + config = PretrainedConfig( + architectures=["BailingMoeV2_5ForCausalLM"], + hidden_size=4096, + kv_lora_rank=512, + num_attention_heads=32, + num_experts=256, + num_hidden_layers=32, + num_key_value_heads=32, + num_nextn_predict_layers=1, + qk_rope_head_dim=64, + vocab_size=157184, + ) + config.model_type = "bailing_hybrid" + return config + + +def test_bailing_hybrid_mtp_hf_config_override(): + config = _bailing_config() + + overridden = SpeculativeConfig.hf_config_override(config) + + assert overridden.model_type == "bailing_hybrid_mtp" + assert overridden.architectures == ["BailingMoeV25MTPModel"] + assert overridden.n_predict == 1 + assert "bailing_hybrid_mtp" in MTPModelTypes.__args__ + + +def test_bailing_hybrid_mtp_model_arch_config(): + config = _bailing_config() + config.model_type = "bailing_hybrid_mtp" + config.architectures = ["BailingMoeV25MTPModel"] + + model_arch_config = BailingHybridMTPModelArchConfigConvertor( + config, config + ).convert() + + assert model_arch_config.model_type == "bailing_hybrid_mtp" + assert model_arch_config.architectures == ["BailingMoeV25MTPModel"] + assert model_arch_config.total_num_hidden_layers == 1 + assert model_arch_config.is_deepseek_mla diff --git a/tests/config/test_speculative_draft_hf_overrides.py b/tests/config/test_speculative_draft_hf_overrides.py new file mode 100644 index 00000000000..ddb8752a80d --- /dev/null +++ b/tests/config/test_speculative_draft_hf_overrides.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for SpeculativeConfig.compose_draft_hf_overrides. + +Callable ``hf_overrides`` on the target model config (e.g. the +``dummy_hf_overrides`` shrink used by ``tests/models/test_initialization.py``) +must also be applied when building the draft ``ModelConfig``. Otherwise a +draft belonging to a large target model is instantiated at full size even +when the target itself is shrunk — which is what kept spec-decode archs like +``EagleMistralLarge3ForCausalLM`` stuck at ``is_available_online=False`` +("TODO: revert once figuring out OOM in CI"). +""" + +import functools + +import pytest +from transformers import PretrainedConfig + +from vllm.config.speculative import SpeculativeConfig + + +def _make_hf_config(**kwargs) -> PretrainedConfig: + defaults = dict( + architectures=["LlamaForCausalLM"], + model_type="llama", + num_hidden_layers=64, + ) + defaults.update(kwargs) + return PretrainedConfig(**defaults) + + +@pytest.mark.cpu_test +def test_dict_overrides_are_not_forwarded_to_draft(): + """Dict overrides are target-specific key patches; the draft must get + only the architecture-mapping override.""" + composed = SpeculativeConfig.compose_draft_hf_overrides( + {"max_position_embeddings": 1234} + ) + assert composed is SpeculativeConfig.hf_config_override + + +@pytest.mark.cpu_test +def test_none_overrides_fall_back_to_arch_mapping(): + composed = SpeculativeConfig.compose_draft_hf_overrides(None) + assert composed is SpeculativeConfig.hf_config_override + + +@pytest.mark.cpu_test +def test_callable_overrides_reach_the_draft_config(): + """A callable override (config-to-config transform) composes with the + architecture-mapping override and is applied to the draft config.""" + + def shrink(hf_config: PretrainedConfig) -> PretrainedConfig: + hf_config.num_hidden_layers = 1 + return hf_config + + composed = SpeculativeConfig.compose_draft_hf_overrides(shrink) + assert composed is not SpeculativeConfig.hf_config_override + + out = composed(_make_hf_config()) + # The shrink transform must have been applied to the draft config. + assert out.num_hidden_layers == 1 + + +@pytest.mark.cpu_test +def test_arch_mapping_applies_before_callable_override(): + """The static arch-mapping override runs first, so the user callable + observes (and may adjust) the post-mapping config.""" + seen_architectures: list[str] = [] + + def record(hf_config: PretrainedConfig) -> PretrainedConfig: + seen_architectures.append(hf_config.architectures[0]) + return hf_config + + composed = SpeculativeConfig.compose_draft_hf_overrides(record) + + # MiMo is one of the arch-mapped model types: hf_config_override + # rewrites architectures to ["MiMoMTPModel"]. + mimo = _make_hf_config( + architectures=["MiMoForCausalLM"], + model_type="mimo", + num_nextn_predict_layers=1, + ) + composed(mimo) + assert seen_architectures == ["MiMoMTPModel"] + + +def _module_level_shrink(hf_config: PretrainedConfig) -> PretrainedConfig: + hf_config.num_hidden_layers = 1 + return hf_config + + +@pytest.mark.cpu_test +def test_composed_override_is_picklable(): + """The draft ``ModelConfig`` is sent to spawned engine-core processes, so + the composed override must be picklable. A nested local closure is not + (it raised ``Can't get local object`` on DFlashDraftModel); a + ``functools.partial`` over a module-referenceable static method is. + Guard against regressing to a closure.""" + composed = SpeculativeConfig.compose_draft_hf_overrides(_module_level_shrink) + + assert isinstance(composed, functools.partial) + assert composed.func is SpeculativeConfig._apply_composed_hf_override + + out = composed(_make_hf_config()) + assert out.num_hidden_layers == 1 diff --git a/tests/distributed/test_rocm_aiter_custom_ar.py b/tests/distributed/test_rocm_aiter_custom_ar.py new file mode 100644 index 00000000000..0b85f36410d --- /dev/null +++ b/tests/distributed/test_rocm_aiter_custom_ar.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import ray +import torch +import torch.distributed as dist + +from vllm._aiter_ops import is_aiter_found, rocm_aiter_ops +from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa +from vllm.distributed.parallel_state import get_tp_group, graph_capture +from vllm.envs import disable_envs_cache +from vllm.platforms import current_platform + +from ..utils import ( + assert_rocm_custom_allreduce_backend_state, + ensure_model_parallel_initialized, + init_test_distributed_environment, + multi_gpu_test, + multi_process_parallel, +) + +pytestmark = pytest.mark.skipif( + not current_platform.is_rocm(), + reason="ROCm-only AITER custom allreduce tests", +) + +test_cases = [ + ((2, 7168), torch.float16), + ((2, 7168), torch.bfloat16), + ((128, 8192), torch.float16), + ((128, 8192), torch.bfloat16), +] + + +def _configure_aiter_custom_ar_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising=False) + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_CUSTOM_AR", "1") + monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", "NONE") + disable_envs_cache() + rocm_aiter_ops.refresh_env_variables() + + +def _assert_aiter_handles_input(inp: torch.Tensor) -> None: + aiter_ar_comm = get_tp_group().device_communicator.aiter_ar_comm + assert aiter_ar_comm is not None + assert aiter_ar_comm.should_custom_ar(inp), ( + f"AITER CustomAllreduce does not support input shape {inp.shape}." + ) + + +@ray.remote(num_gpus=1, max_calls=1) +def graph_allreduce( + monkeypatch: pytest.MonkeyPatch, + tp_size, + pp_size, + rank, + distributed_init_port, +) -> None: + with monkeypatch.context() as m: + _configure_aiter_custom_ar_env(m) + + device = torch.device(f"cuda:{rank}") + torch.accelerator.set_device_index(device) + init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) + ensure_model_parallel_initialized(tp_size, pp_size) + assert_rocm_custom_allreduce_backend_state(True, "NONE") + group = get_tp_group().device_group + + # A small all_reduce for warmup. + # this is needed because device communicators might be created lazily + # (e.g. NCCL). This will ensure that the communicator is initialized + # before any communication happens, so that this group can be used for + # graph capture immediately. + data = torch.zeros(1) + data = data.to(device=device) + dist.all_reduce(data, group=group) + torch.accelerator.synchronize() + del data + + for shape, dtype in test_cases: + with graph_capture(device=device) as graph_capture_context: + inp = torch.ones(shape, dtype=dtype, device=device) + _assert_aiter_handles_input(inp) + expected = inp * tp_size + + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=graph_capture_context.stream): + out = tensor_model_parallel_all_reduce(inp) + + graph.replay() + torch.testing.assert_close(out, expected) + + +@ray.remote(num_gpus=1, max_calls=1) +def eager_allreduce( + monkeypatch: pytest.MonkeyPatch, + tp_size, + pp_size, + rank, + distributed_init_port, +) -> None: + with monkeypatch.context() as m: + _configure_aiter_custom_ar_env(m) + + device = torch.device(f"cuda:{rank}") + torch.accelerator.set_device_index(device) + init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) + ensure_model_parallel_initialized(tp_size, pp_size) + assert_rocm_custom_allreduce_backend_state(True, "NONE") + + for shape, dtype in test_cases: + inp = torch.ones(shape, dtype=dtype, device=device) + _assert_aiter_handles_input(inp) + expected = inp * tp_size + out = tensor_model_parallel_all_reduce(inp) + torch.testing.assert_close(out, expected) + + +@pytest.mark.skipif(not is_aiter_found(), reason="AITER is not installed") +@multi_gpu_test(num_gpus=2) +@pytest.mark.parametrize("tp_size", [2]) +@pytest.mark.parametrize("pipeline_parallel_size", [1]) +@pytest.mark.parametrize("test_target", [eager_allreduce, graph_allreduce]) +def test_rocm_aiter_custom_allreduce( + monkeypatch: pytest.MonkeyPatch, + tp_size, + pipeline_parallel_size, + test_target, +): + multi_process_parallel(monkeypatch, tp_size, pipeline_parallel_size, test_target) diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py index 8a13b24dc52..f3423745ca5 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/tests/distributed/test_weight_transfer.py @@ -191,7 +191,7 @@ class TestNCCLEngineParsing: return NCCLWeightTransferEngine( config, create_mock_vllm_config(), - "cuda", + torch.device("cuda"), MagicMock(spec=torch.nn.Module), ) @@ -240,21 +240,30 @@ class TestEngineRegistry: def test_create_engine_nccl(self): config = WeightTransferConfig(backend="nccl") engine = WeightTransferEngineFactory.create_engine( - config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module) + config, + create_mock_vllm_config(), + torch.device("cuda"), + MagicMock(spec=torch.nn.Module), ) assert isinstance(engine, NCCLWeightTransferEngine) def test_create_engine_ipc(self): config = WeightTransferConfig(backend="ipc") engine = WeightTransferEngineFactory.create_engine( - config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module) + config, + create_mock_vllm_config(), + torch.device("cuda"), + MagicMock(spec=torch.nn.Module), ) assert isinstance(engine, IPCWeightTransferEngine) def test_create_engine_sparse_nccl(self): config = WeightTransferConfig(backend="sparse_nccl") engine = WeightTransferEngineFactory.create_engine( - config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module) + config, + create_mock_vllm_config(), + torch.device("cuda"), + MagicMock(spec=torch.nn.Module), ) assert isinstance(engine, SparseNCCLWeightTransferEngine) @@ -264,7 +273,7 @@ class TestEngineRegistry: WeightTransferEngineFactory.create_engine( config, create_mock_vllm_config(), - "cuda", + torch.device("cuda"), MagicMock(spec=torch.nn.Module), ) @@ -284,7 +293,7 @@ class TestSparseNCCLPatchApplication: def _make_engine(self, model): config = WeightTransferConfig(backend="sparse_nccl") return SparseNCCLWeightTransferEngine( - config, create_mock_vllm_config(), "cpu", model + config, create_mock_vllm_config(), torch.device("cpu"), model ) def _make_model(self, numel: int = 8): @@ -382,7 +391,10 @@ def test_nccl_receive_weights_without_init_raises(): config = WeightTransferConfig(backend="nccl") engine = NCCLWeightTransferEngine( - config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module) + config, + create_mock_vllm_config(), + torch.device("cuda"), + MagicMock(spec=torch.nn.Module), ) update_info = NCCLWeightTransferUpdateInfo( @@ -400,7 +412,10 @@ def test_sparse_nccl_receive_weights_without_init_raises(): config = WeightTransferConfig(backend="sparse_nccl") engine = SparseNCCLWeightTransferEngine( - config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module) + config, + create_mock_vllm_config(), + torch.device("cuda"), + MagicMock(spec=torch.nn.Module), ) update_info = SparseNCCLWeightTransferUpdateInfo( @@ -495,7 +510,9 @@ def inference_receive_tensor( vllm_config.model_config = MagicMock() recorder = Recorder() - engine = NCCLWeightTransferEngine(config, vllm_config, "cuda", recorder) + engine = NCCLWeightTransferEngine( + config, vllm_config, torch.device("cuda"), recorder + ) # Transport-only test: bypass the set_current_vllm_config context that # receive_weights enters, since vllm_config here is a mock. import vllm.config as _vllm_config_mod @@ -664,7 +681,9 @@ def inference_receive_sparse_tensor( num_updates_list=[3], ) - engine = SparseNCCLWeightTransferEngine(config, vllm_config, "cuda", model) + engine = SparseNCCLWeightTransferEngine( + config, vllm_config, torch.device("cuda"), model + ) from vllm.distributed.weight_transfer.nccl_common import ( NCCLWeightTransferInitInfo, ) @@ -879,7 +898,7 @@ class TestIPCEngineParsing: return IPCWeightTransferEngine( config, create_mock_vllm_config(), - "cuda", + torch.device("cuda"), MagicMock(spec=torch.nn.Module), ) @@ -1068,7 +1087,9 @@ def inference_receive_ipc_tensor( vllm_config.model_config = MagicMock() recorder = Recorder() - engine = IPCWeightTransferEngine(config, vllm_config, "cuda", recorder) + engine = IPCWeightTransferEngine( + config, vllm_config, _get_ray_assigned_device(), recorder + ) # Transport-only test: bypass the set_current_vllm_config context that # receive_weights enters, since vllm_config here is a mock. import vllm.config as _vllm_config_mod @@ -1173,7 +1194,10 @@ def test_ipc_receive_weights_missing_gpu_uuid_raises(): config = WeightTransferConfig(backend="ipc") engine = IPCWeightTransferEngine( - config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module) + config, + create_mock_vllm_config(), + torch.device("cuda:0"), + MagicMock(spec=torch.nn.Module), ) dummy_tensor = torch.ones(10, 10, device="cuda:0") diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat.py b/tests/entrypoints/openai/chat_completion/test_serving_chat.py index 7b480bdba67..b43ac31d3dc 100644 --- a/tests/entrypoints/openai/chat_completion/test_serving_chat.py +++ b/tests/entrypoints/openai/chat_completion/test_serving_chat.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio import json +from collections.abc import AsyncIterator from contextlib import suppress from dataclasses import dataclass, field from typing import Any @@ -19,6 +20,7 @@ from tests.entrypoints.openai.utils import ( from tests.utils import RemoteOpenAIServer from vllm._aiter_ops import is_aiter_found_and_supported from vllm.config import MultiModalConfig +from vllm.entrypoints.generate.base.serving import build_per_request_timing_metrics from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ChatCompletionResponse, @@ -50,9 +52,17 @@ from vllm.tokenizers import get_tokenizer from vllm.tokenizers.mistral import MistralTokenizer from vllm.tokenizers.registry import cached_tokenizer_from_config from vllm.v1.engine.async_llm import AsyncLLM +from vllm.v1.metrics.stats import RequestStateStats GPT_OSS_MODEL_NAME = "openai/gpt-oss-20b" GPT_OSS_SPECULATOR_NAME = "RedHatAI/gpt-oss-20b-speculator.eagle3" +_PER_REQUEST_STATS = RequestStateStats( + queued_ts=1.0, + scheduled_ts=1.5, + first_token_ts=2.0, + last_token_ts=3.0, + num_generation_tokens=2, +) @pytest.fixture(scope="module") @@ -608,6 +618,169 @@ def _build_serving_chat( return serving_chat +def _build_minimal_metrics_serving_chat( + enable_per_request_metrics: bool, + enable_force_include_usage: bool = False, +) -> OpenAIServingChat: + serving = OpenAIServingChat.__new__(OpenAIServingChat) + serving.response_role = "assistant" + serving.parser_cls = None + serving.enable_auto_tools = False + serving.enable_prompt_tokens_details = False + serving.enable_log_outputs = False + serving.enable_log_deltas = False + serving.enable_force_include_usage = enable_force_include_usage + serving.request_logger = None + serving.system_fingerprint = None + serving.enable_per_request_metrics = enable_per_request_metrics + return serving + + +def _make_metrics_request_output( + metrics: RequestStateStats | None = _PER_REQUEST_STATS, + token_ids: tuple[int, ...] = (100, 101), +) -> RequestOutput: + return RequestOutput( + request_id="test-id", + prompt="Test prompt", + prompt_token_ids=[1, 2, 3], + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=0, + text="Hello", + token_ids=list(token_ids), + cumulative_logprob=None, + logprobs=None, + finish_reason="stop", + ) + ], + finished=True, + metrics=metrics, + ) + + +async def _single_request_output( + request_output: RequestOutput, +) -> AsyncIterator[RequestOutput]: + yield request_output + + +async def _collect_metrics_stream_chunks( + serving: OpenAIServingChat, + request: ChatCompletionRequest, +) -> list[dict[str, Any]]: + chunks: list[dict[str, Any]] = [] + async for line in serving.chat_completion_stream_generator( + request, + _single_request_output(_make_metrics_request_output()), + "chatcmpl-test-id", + "test-model", + conversation=[{"role": "user", "content": "Test"}], + tokenizer=MagicMock(), + request_metadata=RequestResponseMetadata(request_id="chatcmpl-test-id"), + ): + line = line.strip() + if not line.startswith("data: "): + continue + payload = line[len("data: ") :] + if payload != "[DONE]": + chunks.append(json.loads(payload)) + return chunks + + +def test_build_per_request_timing_metrics_valid_timestamps(): + metrics = build_per_request_timing_metrics( + _PER_REQUEST_STATS, num_generation_tokens=10 + ) + + assert metrics.time_to_first_token_ms == pytest.approx(500.0) + assert metrics.generation_time_ms == pytest.approx(1000.0) + assert metrics.queue_time_ms == pytest.approx(500.0) + assert metrics.mean_itl_ms == pytest.approx(1000.0 / 9, rel=1e-4) + assert metrics.tokens_per_second == pytest.approx(10.0 / 1.5, rel=1e-4) + + +@pytest.mark.asyncio +async def test_chat_per_request_metrics_follow_server_flag(): + request = ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "Test prompt"}], + max_tokens=10, + stream=False, + ) + request_output = _make_metrics_request_output() + + disabled_serving = _build_minimal_metrics_serving_chat( + enable_per_request_metrics=False + ) + disabled_response = await disabled_serving.chat_completion_full_generator( + request, + _single_request_output(request_output), + "chatcmpl-test-id", + "test-model", + conversation=[{"role": "user", "content": "Test"}], + tokenizer=MagicMock(), + request_metadata=RequestResponseMetadata(request_id="chatcmpl-test-id"), + ) + assert disabled_response.metrics is None + + enabled_serving = _build_minimal_metrics_serving_chat( + enable_per_request_metrics=True + ) + enabled_response = await enabled_serving.chat_completion_full_generator( + request, + _single_request_output(request_output), + "chatcmpl-test-id", + "test-model", + conversation=[{"role": "user", "content": "Test"}], + tokenizer=MagicMock(), + request_metadata=RequestResponseMetadata(request_id="chatcmpl-test-id"), + ) + assert enabled_response.metrics is not None + assert enabled_response.metrics.time_to_first_token_ms == pytest.approx(500.0) + + +@pytest.mark.asyncio +async def test_chat_per_request_metrics_suppressed_for_n_greater_than_one(): + serving = _build_minimal_metrics_serving_chat(enable_per_request_metrics=True) + response = await serving.chat_completion_full_generator( + ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "Test prompt"}], + max_tokens=10, + stream=False, + n=2, + ), + _single_request_output(_make_metrics_request_output()), + "chatcmpl-test-id", + "test-model", + conversation=[{"role": "user", "content": "Test"}], + tokenizer=MagicMock(), + request_metadata=RequestResponseMetadata(request_id="chatcmpl-test-id"), + ) + assert response.metrics is None + + +@pytest.mark.asyncio +async def test_chat_streaming_metrics_ride_on_usage_chunk(): + serving = _build_minimal_metrics_serving_chat(enable_per_request_metrics=True) + chunks = await _collect_metrics_stream_chunks( + serving, + ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "Test prompt"}], + max_tokens=10, + stream=True, + stream_options={"include_usage": True}, + ), + ) + + usage_chunks = [chunk for chunk in chunks if chunk.get("usage")] + assert usage_chunks + assert usage_chunks[-1]["metrics"]["time_to_first_token_ms"] == pytest.approx(500.0) + + @dataclass class MockEngine: model_config: MockModelConfig = field(default_factory=MockModelConfig) diff --git a/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py b/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py index 3e9e4850f07..4c574673817 100644 --- a/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py +++ b/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py @@ -183,27 +183,39 @@ async def test_thinking_token_budget_mixed_requests(client: openai.AsyncOpenAI): async def test_thinking_token_budget_limits_reasoning(client: openai.AsyncOpenAI): """Test that thinking_token_budget limits the number of reasoning tokens. - Counts non-empty streaming ``delta.reasoning`` chunks (coarse proxy; each - chunk may represent multiple decode tokens — see - ``_count_reasoning_decode_token_ids_between_markers`` and the Qwen3.5 MTP - test for id-based checks). + Counts reasoning decode tokens by id, which is robust to how tokens are + grouped into streamed chunks (a single chunk can carry several tokens under + async scheduling / stream_interval > 1). Counting chunks under-counts. """ - reasoning_token_count = 0 + tokenizer = get_tokenizer(tokenizer_name=MODEL_NAME) + start_ids = list(tokenizer.encode(REASONING_START_STR, add_special_tokens=False)) + end_ids = list(tokenizer.encode(REASONING_END_STR, add_special_tokens=False)) + + prompt_token_ids: list[int] = [] + decode_token_ids: list[int] = [] stream = await client.chat.completions.create( model=MODEL_NAME, messages=MESSAGES, max_tokens=100, stream=True, - extra_body={"thinking_token_budget": THINK_BUDGET}, + extra_body={"thinking_token_budget": THINK_BUDGET, "return_token_ids": True}, ) async for chunk in stream: - delta = chunk.choices[0].delta - if getattr(delta, "reasoning", None): - reasoning_token_count += 1 + if not chunk.choices: + continue + if getattr(chunk, "prompt_token_ids", None): + prompt_token_ids = list(chunk.prompt_token_ids) + delta_ids = getattr(chunk.choices[0], "token_ids", None) + if delta_ids: + decode_token_ids.extend(delta_ids) + reasoning_token_count = _count_reasoning_decode_token_ids_between_markers( + prompt_token_ids + decode_token_ids, start_ids, end_ids + ) + assert reasoning_token_count is not None, "missing reasoning start marker in ids" assert reasoning_token_count == THINK_BUDGET, ( - f"reasoning tokens ({reasoning_token_count}) exceeded " + f"reasoning tokens ({reasoning_token_count}) != " f"thinking_token_budget ({THINK_BUDGET})" ) diff --git a/tests/entrypoints/openai/completion/test_completion_error.py b/tests/entrypoints/openai/completion/test_completion_error.py index 062c3e7583a..5906063c54a 100644 --- a/tests/entrypoints/openai/completion/test_completion_error.py +++ b/tests/entrypoints/openai/completion/test_completion_error.py @@ -11,7 +11,10 @@ from pydantic import ValidationError from vllm.config.multimodal import MultiModalConfig from vllm.entrypoints.openai.completion.protocol import CompletionRequest from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion -from vllm.entrypoints.openai.engine.protocol import GenerationError +from vllm.entrypoints.openai.engine.protocol import ( + GenerationError, + RequestResponseMetadata, +) from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.scale_out.render.serving import ServingRender @@ -20,9 +23,17 @@ from vllm.renderers.hf import HfRenderer from vllm.renderers.online_renderer import OnlineRenderer from vllm.tokenizers.registry import cached_tokenizer_from_config from vllm.v1.engine.async_llm import AsyncLLM +from vllm.v1.metrics.stats import RequestStateStats MODEL_NAME = "openai-community/gpt2" MODEL_NAME_SHORT = "gpt2" +_PER_REQUEST_STATS = RequestStateStats( + queued_ts=1.0, + scheduled_ts=1.5, + first_token_ts=2.0, + last_token_ts=3.0, + num_generation_tokens=2, +) BASE_MODEL_PATHS = [ BaseModelPath(name=MODEL_NAME, model_path=MODEL_NAME), BaseModelPath(name=MODEL_NAME_SHORT, model_path=MODEL_NAME_SHORT), @@ -93,6 +104,39 @@ def _build_serving_completion(engine: AsyncLLM) -> OpenAIServingCompletion: ) +def _build_minimal_metrics_serving_completion( + enable_per_request_metrics: bool, +) -> OpenAIServingCompletion: + serving = OpenAIServingCompletion.__new__(OpenAIServingCompletion) + serving.enable_prompt_tokens_details = False + serving.system_fingerprint = None + serving.enable_per_request_metrics = enable_per_request_metrics + return serving + + +def _make_metrics_request_output( + metrics: RequestStateStats | None = _PER_REQUEST_STATS, +) -> RequestOutput: + return RequestOutput( + request_id="test-id", + prompt="Test prompt", + prompt_token_ids=[1, 2, 3], + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=0, + text="Hello", + token_ids=[100, 101], + cumulative_logprob=None, + logprobs=None, + finish_reason="stop", + ) + ], + finished=True, + metrics=metrics, + ) + + def _build_renderer(model_config: MockModelConfig): return HfRenderer( MockVllmConfig(model_config, parallel_config=MockParallelConfig()), @@ -100,6 +144,58 @@ def _build_renderer(model_config: MockModelConfig): ) +def test_completion_per_request_metrics_follow_server_flag(): + request = CompletionRequest(model=MODEL_NAME, prompt="Test prompt", max_tokens=10) + request_output = _make_metrics_request_output() + + disabled_serving = _build_minimal_metrics_serving_completion( + enable_per_request_metrics=False + ) + disabled_response = disabled_serving.request_output_to_completion_response( + [request_output], + request, + "cmpl-test-id", + 0, + MODEL_NAME, + None, + RequestResponseMetadata(request_id="cmpl-test-id"), + ) + assert disabled_response.metrics is None + + enabled_serving = _build_minimal_metrics_serving_completion( + enable_per_request_metrics=True + ) + enabled_response = enabled_serving.request_output_to_completion_response( + [request_output], + request, + "cmpl-test-id", + 0, + MODEL_NAME, + None, + RequestResponseMetadata(request_id="cmpl-test-id"), + ) + assert enabled_response.metrics is not None + assert enabled_response.metrics.time_to_first_token_ms == pytest.approx(500.0) + + +def test_completion_per_request_metrics_suppressed_for_multiple_prompts(): + serving = _build_minimal_metrics_serving_completion(enable_per_request_metrics=True) + response = serving.request_output_to_completion_response( + [_make_metrics_request_output(), _make_metrics_request_output()], + CompletionRequest( + model=MODEL_NAME, + prompt=["Test prompt", "Another prompt"], + max_tokens=10, + ), + "cmpl-test-id", + 0, + MODEL_NAME, + None, + RequestResponseMetadata(request_id="cmpl-test-id"), + ) + assert response.metrics is None + + @pytest.mark.asyncio async def test_completion_error_non_stream(): """test finish_reason='error' returns 500 InternalServerError (non-streaming)""" diff --git a/tests/entrypoints/openai/correctness/test_lmeval.py b/tests/entrypoints/openai/correctness/test_lmeval.py index 5b23b423902..aad1b5e0624 100644 --- a/tests/entrypoints/openai/correctness/test_lmeval.py +++ b/tests/entrypoints/openai/correctness/test_lmeval.py @@ -71,8 +71,9 @@ def test_lm_eval_accuracy_v1_engine(): more_args = [] - # Limit compilation time for V1 - if current_platform.is_tpu(): + # Limit compilation time for V1 on TPU + # Avoid OOM on XPU + if current_platform.is_tpu() or current_platform.is_xpu(): more_args = ["--max-num-seqs", "64"] run_test(more_args) diff --git a/tests/entrypoints/openai/responses/test_namespace_tool_separator.py b/tests/entrypoints/openai/responses/test_namespace_tool_separator.py new file mode 100644 index 00000000000..c895092b45f --- /dev/null +++ b/tests/entrypoints/openai/responses/test_namespace_tool_separator.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json + +import openai # use the official client for correctness check +import pytest + +MODEL_NAME = "Qwen/Qwen3-1.7B" +NAMESPACE = "mcp__computer_use" +TOOL_NAME = "get_app_state" +FLAT_TOOL_NAME = f"{NAMESPACE}__{TOOL_NAME}" + +tools = [ + { + "type": "namespace", + "name": NAMESPACE, + "description": "Computer control tools.", + "tools": [ + { + "type": "function", + "name": TOOL_NAME, + "description": "Get the current state of a desktop application.", + "parameters": { + "type": "object", + "properties": { + "app": { + "type": "string", + "description": "Application name, for example Chrome.", + } + }, + "required": ["app"], + "additionalProperties": False, + }, + } + ], + } +] + +prompt = [ + { + "role": "user", + "content": "Use the computer app state tool to inspect Google Chrome.", + }, +] + + +def _assert_namespace_tool_call(tool_call) -> None: + assert tool_call.type == "function_call" + assert tool_call.name == TOOL_NAME + assert tool_call.namespace == NAMESPACE + assert tool_call.name != FLAT_TOOL_NAME + + args = json.loads(tool_call.arguments) + assert args["app"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model_name", [MODEL_NAME]) +async def test_namespace_tool_separator(client: openai.AsyncOpenAI, model_name: str): + response = await client.responses.create( + model=model_name, + input=prompt, + tools=tools, + temperature=0.0, + ) + + assert len(response.output) >= 1 + tool_call = next( + (out for out in response.output if out.type == "function_call"), None + ) + assert tool_call is not None + _assert_namespace_tool_call(tool_call) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model_name", [MODEL_NAME]) +async def test_namespace_tool_separator_streaming( + client: openai.AsyncOpenAI, model_name: str +): + stream = await client.responses.create( + model=model_name, + input=prompt, + tools=tools, + temperature=0.0, + stream=True, + ) + events = [event async for event in stream] + + added_call = next( + ( + event.item + for event in events + if event.type == "response.output_item.added" + and getattr(event.item, "type", None) == "function_call" + ), + None, + ) + done_call = next( + ( + event.item + for event in events + if event.type == "response.output_item.done" + and getattr(event.item, "type", None) == "function_call" + ), + None, + ) + + assert added_call is not None + assert added_call.name == TOOL_NAME + assert added_call.namespace == NAMESPACE + + assert done_call is not None + _assert_namespace_tool_call(done_call) diff --git a/tests/entrypoints/openai/test_cli_args.py b/tests/entrypoints/openai/test_cli_args.py index 58dd328b325..1f764202e55 100644 --- a/tests/entrypoints/openai/test_cli_args.py +++ b/tests/entrypoints/openai/test_cli_args.py @@ -206,6 +206,14 @@ def test_chat_template_validation_for_sad_paths(serve_parser): validate_parsed_serve_args(args) +def test_per_request_metrics_requires_log_stats(serve_parser): + args = serve_parser.parse_args( + args=["--enable-per-request-metrics", "--disable-log-stats"] + ) + with pytest.raises(ValueError): + validate_parsed_serve_args(args) + + @pytest.mark.parametrize( "cli_args, expected_middleware", [ diff --git a/tests/evals/gsm8k/configs/humming/Qwen2-1.5B-Instruct-FP8W8-humming-act-fp8.yaml b/tests/evals/gsm8k/configs/humming/Qwen2-1.5B-Instruct-FP8W8-humming-act-fp8.yaml new file mode 100644 index 00000000000..ba292eb9724 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen2-1.5B-Instruct-FP8W8-humming-act-fp8.yaml @@ -0,0 +1,10 @@ +model_name: "nm-testing/Qwen2-1.5B-Instruct-FP8W8" +accuracy_threshold: 0.55 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --linear-backend humming +env: + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen2-1.5B-Instruct-FP8W8-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen2-1.5B-Instruct-FP8W8-humming.yaml new file mode 100644 index 00000000000..3179c3251b2 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen2-1.5B-Instruct-FP8W8-humming.yaml @@ -0,0 +1,8 @@ +model_name: "nm-testing/Qwen2-1.5B-Instruct-FP8W8" +accuracy_threshold: 0.55 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --linear-backend humming diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-0.6B-MXFP8-humming-act-fp8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-0.6B-MXFP8-humming-act-fp8.yaml new file mode 100644 index 00000000000..66888312ef0 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-0.6B-MXFP8-humming-act-fp8.yaml @@ -0,0 +1,10 @@ +model_name: "mgoin/Qwen3-0.6B-MXFP8" +accuracy_threshold: 0.39 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --linear-backend humming +env: + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-0.6B-MXFP8-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-0.6B-MXFP8-humming.yaml new file mode 100644 index 00000000000..b83d9e6a9e9 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-0.6B-MXFP8-humming.yaml @@ -0,0 +1,8 @@ +model_name: "mgoin/Qwen3-0.6B-MXFP8" +accuracy_threshold: 0.39 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --linear-backend humming diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming-act-fp8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming-act-fp8.yaml new file mode 100644 index 00000000000..d3de5b3792f --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming-act-fp8.yaml @@ -0,0 +1,13 @@ +model_name: "QuixiAI/Qwen3-30B-A3B-AWQ" +accuracy_threshold: 0.90 +num_questions: 1319 +num_fewshot: 5 +gen_prefix: " \n\n\n" +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming + --dtype bfloat16 +env: + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming-act-int8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming-act-int8.yaml new file mode 100644 index 00000000000..6b76efc4e22 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming-act-int8.yaml @@ -0,0 +1,13 @@ +model_name: "QuixiAI/Qwen3-30B-A3B-AWQ" +accuracy_threshold: 0.90 +num_questions: 1319 +num_fewshot: 5 +gen_prefix: " \n\n\n" +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming + --dtype bfloat16 +env: + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"int8"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming.yaml new file mode 100644 index 00000000000..310255f6bf4 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming.yaml @@ -0,0 +1,11 @@ +model_name: "QuixiAI/Qwen3-30B-A3B-AWQ" +accuracy_threshold: 0.90 +num_questions: 1319 +num_fewshot: 5 +gen_prefix: " \n\n\n" +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming + --dtype bfloat16 diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-act-fp8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-act-fp8.yaml new file mode 100644 index 00000000000..6d2702d8f4c --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-act-fp8.yaml @@ -0,0 +1,11 @@ +model_name: "nm-testing/Qwen3-30B-A3B-FP8-block" +accuracy_threshold: 0.86 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming +env: + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming.yaml new file mode 100644 index 00000000000..8ca4777ce5e --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming.yaml @@ -0,0 +1,9 @@ +model_name: "nm-testing/Qwen3-30B-A3B-FP8-block" +accuracy_threshold: 0.86 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-Fp8-v1-humming-act-fp8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-Fp8-v1-humming-act-fp8.yaml new file mode 100644 index 00000000000..618e7fdcc35 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-Fp8-v1-humming-act-fp8.yaml @@ -0,0 +1,11 @@ +model_name: "nm-testing/Qwen3-30B-A3B-Fp8-v1" +accuracy_threshold: 0.86 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming +env: + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-Fp8-v1-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-Fp8-v1-humming.yaml new file mode 100644 index 00000000000..71aabcef99d --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-Fp8-v1-humming.yaml @@ -0,0 +1,9 @@ +model_name: "nm-testing/Qwen3-30B-A3B-Fp8-v1" +accuracy_threshold: 0.86 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-GPTQ-Int4-humming-act-fp8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-GPTQ-Int4-humming-act-fp8.yaml new file mode 100644 index 00000000000..251fb252cf8 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-GPTQ-Int4-humming-act-fp8.yaml @@ -0,0 +1,11 @@ +model_name: "Qwen/Qwen3-30B-A3B-GPTQ-Int4" +accuracy_threshold: 0.86 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming +env: + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-GPTQ-Int4-humming-act-int8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-GPTQ-Int4-humming-act-int8.yaml new file mode 100644 index 00000000000..d813e05d6f9 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-GPTQ-Int4-humming-act-int8.yaml @@ -0,0 +1,11 @@ +model_name: "Qwen/Qwen3-30B-A3B-GPTQ-Int4" +accuracy_threshold: 0.86 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming +env: + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"int8"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-GPTQ-Int4-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-GPTQ-Int4-humming.yaml new file mode 100644 index 00000000000..ab89ebf5154 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-GPTQ-Int4-humming.yaml @@ -0,0 +1,9 @@ +model_name: "Qwen/Qwen3-30B-A3B-GPTQ-Int4" +accuracy_threshold: 0.86 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming-act-int8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming-act-int8.yaml new file mode 100644 index 00000000000..f77ee1173e1 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming-act-int8.yaml @@ -0,0 +1,11 @@ +model_name: "RedHatAI/Qwen3-30B-A3B-Instruct-2507-quantized.w8a8" +accuracy_threshold: 0.86 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming +env: + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"int8"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming.yaml new file mode 100644 index 00000000000..7b4f82c8458 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming.yaml @@ -0,0 +1,9 @@ +model_name: "RedHatAI/Qwen3-30B-A3B-Instruct-2507-quantized.w8a8" +accuracy_threshold: 0.86 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming-act-fp8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming-act-fp8.yaml index 9b77af67327..d98f91e1f99 100644 --- a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming-act-fp8.yaml +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming-act-fp8.yaml @@ -5,8 +5,7 @@ num_fewshot: 5 server_args: >- --enforce-eager --max-model-len 8192 - --tensor-parallel-size 1 --quantization humming - --kernel-config.enable_flashinfer_autotune=False + --linear-backend humming env: VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming.yaml index 0b1599ff94b..67725fce64a 100644 --- a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming.yaml +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming.yaml @@ -5,6 +5,5 @@ num_fewshot: 5 server_args: >- --enforce-eager --max-model-len 8192 - --tensor-parallel-size 1 --quantization humming - --kernel-config.enable_flashinfer_autotune=False + --linear-backend humming diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-NVFP4-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-NVFP4-humming.yaml new file mode 100644 index 00000000000..a20c1433103 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-NVFP4-humming.yaml @@ -0,0 +1,9 @@ +model_name: "nvidia/Qwen3-30B-A3B-NVFP4" +accuracy_threshold: 0.86 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-fp8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-fp8.yaml new file mode 100644 index 00000000000..b58c0d710e4 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-fp8.yaml @@ -0,0 +1,13 @@ +model_name: "Qwen/Qwen3-30B-A3B" +accuracy_threshold: 0.80 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --tensor-parallel-size 1 + --quantization humming + --kernel-config.enable_flashinfer_autotune=False +env: + VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int5","hadamard_block_size":-1}' + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-int8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-int8.yaml new file mode 100644 index 00000000000..c932091dbe3 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-int8.yaml @@ -0,0 +1,13 @@ +model_name: "Qwen/Qwen3-30B-A3B" +accuracy_threshold: 0.80 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --tensor-parallel-size 1 + --quantization humming + --kernel-config.enable_flashinfer_autotune=False +env: + VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int5","hadamard_block_size":-1}' + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"int8"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming.yaml new file mode 100644 index 00000000000..fa5095cc882 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming.yaml @@ -0,0 +1,12 @@ +model_name: "Qwen/Qwen3-30B-A3B" +accuracy_threshold: 0.80 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --tensor-parallel-size 1 + --quantization humming + --kernel-config.enable_flashinfer_autotune=False +env: + VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int5","hadamard_block_size":-1}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-FP8-humming-act-fp8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-FP8-humming-act-fp8.yaml new file mode 100644 index 00000000000..3af5c03a245 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-FP8-humming-act-fp8.yaml @@ -0,0 +1,12 @@ +model_name: "Qwen/Qwen3.5-35B-A3B-FP8" +accuracy_threshold: 0.90 +num_questions: 1319 +num_fewshot: 5 +gen_prefix: " \n\n\n" +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming +env: + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-FP8-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-FP8-humming.yaml new file mode 100644 index 00000000000..85fce244200 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-FP8-humming.yaml @@ -0,0 +1,10 @@ +model_name: "Qwen/Qwen3.5-35B-A3B-FP8" +accuracy_threshold: 0.90 +num_questions: 1319 +num_fewshot: 5 +gen_prefix: " \n\n\n" +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming diff --git a/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-experts-int8-humming-act-int8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-experts-int8-humming-act-int8.yaml new file mode 100644 index 00000000000..6a85ab384f2 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-experts-int8-humming-act-int8.yaml @@ -0,0 +1,12 @@ +model_name: "Qwen/Qwen3.5-35B-A3B" +accuracy_threshold: 0.90 +num_questions: 1319 +num_fewshot: 5 +gen_prefix: " \n\n\n" +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --quantization experts_int8 +env: + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"int8"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-experts-int8-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-experts-int8-humming.yaml new file mode 100644 index 00000000000..d282fdc7019 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-experts-int8-humming.yaml @@ -0,0 +1,10 @@ +model_name: "Qwen/Qwen3.5-35B-A3B" +accuracy_threshold: 0.90 +num_questions: 1319 +num_fewshot: 5 +gen_prefix: " \n\n\n" +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --quantization experts_int8 diff --git a/tests/evals/gsm8k/configs/humming/Qwen3.5-4B-quantized.w4a16-humming-act-fp8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.5-4B-quantized.w4a16-humming-act-fp8.yaml new file mode 100644 index 00000000000..2a118098fc0 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3.5-4B-quantized.w4a16-humming-act-fp8.yaml @@ -0,0 +1,10 @@ +model_name: "RedHatAI/Qwen3.5-4B-quantized.w4a16" +accuracy_threshold: 0.82 +num_questions: 1319 +num_fewshot: 5 +gen_prefix: " \n\n\n" +server_args: >- + --enforce-eager + --max-model-len 8192 +env: + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3.5-4B-quantized.w4a16-humming-act-int8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.5-4B-quantized.w4a16-humming-act-int8.yaml new file mode 100644 index 00000000000..34f17a4b055 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3.5-4B-quantized.w4a16-humming-act-int8.yaml @@ -0,0 +1,11 @@ +model_name: "RedHatAI/Qwen3.5-4B-quantized.w4a16" +accuracy_threshold: 0.82 +num_questions: 1319 +num_fewshot: 5 +gen_prefix: " \n\n\n" +server_args: >- + --enforce-eager + --max-model-len 8192 + --linear-backend humming +env: + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"int8"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3.5-4B-quantized.w4a16-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.5-4B-quantized.w4a16-humming.yaml new file mode 100644 index 00000000000..b617c61eb3e --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3.5-4B-quantized.w4a16-humming.yaml @@ -0,0 +1,9 @@ +model_name: "RedHatAI/Qwen3.5-4B-quantized.w4a16" +accuracy_threshold: 0.82 +num_questions: 1319 +num_fewshot: 5 +gen_prefix: " \n\n\n" +server_args: >- + --enforce-eager + --max-model-len 8192 + --linear-backend humming diff --git a/tests/evals/gsm8k/configs/humming/Qwen3.6-35B-A3B-NVFP4-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.6-35B-A3B-NVFP4-humming.yaml new file mode 100644 index 00000000000..502ab776f40 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3.6-35B-A3B-NVFP4-humming.yaml @@ -0,0 +1,10 @@ +model_name: "RedHatAI/Qwen3.6-35B-A3B-NVFP4" +accuracy_threshold: 0.91 +num_questions: 1319 +num_fewshot: 5 +gen_prefix: " \n\n\n" +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming diff --git a/tests/evals/gsm8k/configs/humming/config-act-fp8.txt b/tests/evals/gsm8k/configs/humming/config-act-fp8.txt index 05fb6a15838..42ff6be00ef 100644 --- a/tests/evals/gsm8k/configs/humming/config-act-fp8.txt +++ b/tests/evals/gsm8k/configs/humming/config-act-fp8.txt @@ -1,2 +1,9 @@ gpt-oss-20b-humming-act-fp8.yaml Qwen3-30B-A3B-MXFP4A16-humming-act-fp8.yaml +Qwen2-1.5B-Instruct-FP8W8-humming-act-fp8.yaml +Qwen3-0.6B-MXFP8-humming-act-fp8.yaml +Qwen3-30B-A3B-Fp8-v1-humming-act-fp8.yaml +Qwen3-30B-A3B-FP8-block-humming-act-fp8.yaml +Qwen3-30B-A3B-GPTQ-Int4-humming-act-fp8.yaml +Qwen3-30B-A3B-AWQ-humming.yaml +Qwen3.5-35B-A3B-FP8-humming-act-fp8.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-act-int8.txt b/tests/evals/gsm8k/configs/humming/config-act-int8.txt new file mode 100644 index 00000000000..b018b35cbfd --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/config-act-int8.txt @@ -0,0 +1,4 @@ +Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming-act-int8.yaml +Qwen3-30B-A3B-GPTQ-Int4-humming-act-int8.yaml +Qwen3-30B-A3B-AWQ-humming.yaml +Qwen3.5-35B-A3B-experts-int8-humming-act-int8.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-int5wc-hadamard.txt b/tests/evals/gsm8k/configs/humming/config-int5wc-hadamard.txt new file mode 100644 index 00000000000..2c10777a095 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/config-int5wc-hadamard.txt @@ -0,0 +1,3 @@ +Qwen3-30B-A3B-int5wc-hadamard-humming.yaml +Qwen3-30B-A3B-int5wc-hadamard-humming-act-fp8.yaml +Qwen3-30B-A3B-int5wc-hadamard-humming-act-int8.yaml diff --git a/tests/evals/gsm8k/configs/humming/config.txt b/tests/evals/gsm8k/configs/humming/config.txt index 821025365c7..144ed959935 100644 --- a/tests/evals/gsm8k/configs/humming/config.txt +++ b/tests/evals/gsm8k/configs/humming/config.txt @@ -1,2 +1,13 @@ gpt-oss-20b-humming.yaml Qwen3-30B-A3B-MXFP4A16-humming.yaml +Qwen2-1.5B-Instruct-FP8W8-humming.yaml +Qwen3-0.6B-MXFP8-humming.yaml +Qwen3.6-35B-A3B-NVFP4-humming.yaml +Qwen3-30B-A3B-Fp8-v1-humming.yaml +Qwen3-30B-A3B-FP8-block-humming.yaml +Qwen3-30B-A3B-GPTQ-Int4-humming.yaml +Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming.yaml +Qwen3-30B-A3B-NVFP4-humming.yaml +Qwen3-30B-A3B-AWQ-humming.yaml +Qwen3.5-35B-A3B-FP8-humming.yaml +Qwen3.5-35B-A3B-experts-int8-humming.yaml diff --git a/tests/evals/gsm8k/configs/humming/gpt-oss-20b-humming-act-fp8.yaml b/tests/evals/gsm8k/configs/humming/gpt-oss-20b-humming-act-fp8.yaml index 00ba9eccfda..8e0d9535030 100644 --- a/tests/evals/gsm8k/configs/humming/gpt-oss-20b-humming-act-fp8.yaml +++ b/tests/evals/gsm8k/configs/humming/gpt-oss-20b-humming-act-fp8.yaml @@ -5,7 +5,6 @@ num_fewshot: 5 server_args: >- --enforce-eager --max-model-len 8192 - --tensor-parallel-size 1 --moe-backend humming env: VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}' diff --git a/tests/evals/gsm8k/configs/humming/gpt-oss-20b-humming.yaml b/tests/evals/gsm8k/configs/humming/gpt-oss-20b-humming.yaml index e2beb3739b1..7e9b6508a20 100644 --- a/tests/evals/gsm8k/configs/humming/gpt-oss-20b-humming.yaml +++ b/tests/evals/gsm8k/configs/humming/gpt-oss-20b-humming.yaml @@ -5,5 +5,4 @@ num_fewshot: 5 server_args: >- --enforce-eager --max-model-len 8192 - --tensor-parallel-size 1 --moe-backend humming diff --git a/tests/evals/gsm8k/gsm8k_eval.py b/tests/evals/gsm8k/gsm8k_eval.py index 89d91cb4a98..87235a9c378 100644 --- a/tests/evals/gsm8k/gsm8k_eval.py +++ b/tests/evals/gsm8k/gsm8k_eval.py @@ -146,6 +146,7 @@ async def call_vllm_chat_api( def _build_gsm8k_prompts( num_questions: int = 1319, num_shots: int = 5, + gen_prefix: str = "", ) -> tuple[list[str], list[int]]: """Build few-shot GSM8K completion prompts and ground-truth labels.""" if num_questions == 0: @@ -157,14 +158,15 @@ def _build_gsm8k_prompts( for i in range(num_shots): few_shot_examples += ( f"Question: {train_data[i]['question']}\n" - f"Answer: {train_data[i]['answer']}\n\n" + f"Answer:{gen_prefix} {train_data[i]['answer']}\n\n" ) prompts = [] labels = [] for i in range(num_questions): prompts.append( - few_shot_examples + f"Question: {test_data[i]['question']}\nAnswer:" + few_shot_examples + + f"Question: {test_data[i]['question']}\nAnswer:{gen_prefix}" ) labels.append(get_answer_value(test_data[i]["answer"])) @@ -213,6 +215,7 @@ def evaluate_gsm8k( temperature: float = 0.0, seed: int | None = 42, request_timeout_seconds: float = 600, + gen_prefix: str = "", ) -> dict[str, float | int]: """ Evaluate GSM8K accuracy using vLLM serve endpoint. @@ -220,7 +223,7 @@ def evaluate_gsm8k( Returns dict with accuracy, invalid_rate, latency, etc. """ base_url = f"{host}:{port}" - prompts, labels = _build_gsm8k_prompts(num_questions, num_shots) + prompts, labels = _build_gsm8k_prompts(num_questions, num_shots, gen_prefix) num_questions = len(prompts) async def run_async_evaluation(): @@ -278,6 +281,7 @@ def evaluate_gsm8k_offline( num_shots: int = 5, max_tokens: int = 256, temperature: float = 0.0, + gen_prefix: str = "", ) -> dict[str, float | int]: """Evaluate GSM8K accuracy using an offline vllm.LLM object. @@ -286,7 +290,7 @@ def evaluate_gsm8k_offline( """ from vllm import SamplingParams - prompts, labels = _build_gsm8k_prompts(num_questions, num_shots) + prompts, labels = _build_gsm8k_prompts(num_questions, num_shots, gen_prefix) sampling_params = SamplingParams( temperature=temperature, diff --git a/tests/evals/gsm8k/test_gsm8k_correctness.py b/tests/evals/gsm8k/test_gsm8k_correctness.py index c9ec5ff66e5..0c48af6d3c6 100644 --- a/tests/evals/gsm8k/test_gsm8k_correctness.py +++ b/tests/evals/gsm8k/test_gsm8k_correctness.py @@ -70,6 +70,7 @@ def run_gsm8k_eval(eval_config: dict, server_url: str) -> dict: host=host, port=port, request_timeout_seconds=request_timeout_seconds, + gen_prefix=eval_config.get("gen_prefix", ""), ) return results diff --git a/tests/kernels/helion/test_silu_and_mul_per_block_quant.py b/tests/kernels/helion/test_silu_and_mul_per_block_quant.py new file mode 100644 index 00000000000..b8fcd9c8a67 --- /dev/null +++ b/tests/kernels/helion/test_silu_and_mul_per_block_quant.py @@ -0,0 +1,224 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the silu_and_mul_per_block_quant helion kernel +Run `pytest tests/kernels/helion/test_silu_and_mul_per_block_quant.py`. +""" + +from typing import Any + +import pytest +import torch +from torch._subclasses.fake_tensor import FakeTensorMode + +from tests.kernels.helion.utils import skip_if_platform_unsupported +from tests.kernels.quant_utils import FP8_DTYPE +from vllm.kernels.helion.case_key import CaseKey +from vllm.kernels.helion.config_manager import ConfigManager +from vllm.kernels.helion.ops.silu_and_mul_per_block_quant import ( + _pick_cache, + baseline, + pick_config, + silu_and_mul_per_block_quant, +) +from vllm.platforms import current_platform +from vllm.utils.import_utils import has_helion +from vllm.utils.torch_utils import set_random_seed + +if not has_helion(): + pytest.skip( + "Helion is not installed. Install with: pip install vllm[helion]", + allow_module_level=True, + ) + + +def _generate_fake_input( + num_tokens: int, intermediate_size: int, group_size: int +) -> tuple[Any, ...]: + with FakeTensorMode(): + in_dtype: torch.dtype = torch.bfloat16 + out_dtype: torch.dtype = current_platform.fp8_dtype() + scale_dtype: torch.dtype = torch.float32 + input = torch.randn( + num_tokens, 2 * intermediate_size, device="cuda", dtype=in_dtype + ) + result = torch.empty( + num_tokens, intermediate_size, device=input.device, dtype=out_dtype + ) + scale = torch.empty( + (num_tokens, intermediate_size // group_size), + device=input.device, + dtype=scale_dtype, + ) + scale_ub = torch.mean(input).to(scale_dtype) + args = ( + result, + input, + scale, + group_size, + scale_ub, + False, + ) + return args + + +class TestSiluAndMulPerBlockQuantConfigPicker: + def setup_method(self): + _pick_cache.clear() + + def test_config_picker_exact_match(self): + config_keys = [ + CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 16}), + CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 16}), + ] + + args = _generate_fake_input(16, 4096, 128) + selected_key = pick_config(args, config_keys) + assert selected_key == CaseKey( + {"intermediate_size": 4096, "group_size": 128, "num_tokens": 16} + ) + + def test_config_picker_closest_match(self): + config_keys = [ + CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 16}), + CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 32}), + CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 16}), + CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 32}), + CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 16}), + CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 32}), + CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 16}), + CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 32}), + ] + + args = _generate_fake_input(20, 3000, 70) + selected_key = pick_config(args, config_keys) + assert selected_key == CaseKey( + {"intermediate_size": 2048, "group_size": 64, "num_tokens": 32} + ) + + def test_config_picker_no_configs(self): + config_keys: list[dict] = [] + + args = _generate_fake_input(16, 4096, 128) + selected_key = pick_config(args, config_keys) + assert selected_key is None + + def test_config_picker_fallback_to_largest(self): + config_keys = [ + CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 16}), + CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 32}), + CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 16}), + CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 32}), + CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 16}), + CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 32}), + CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 16}), + CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 32}), + ] + + args = _generate_fake_input(64, 8192, 256) + selected_key = pick_config(args, config_keys) + assert selected_key == CaseKey( + {"intermediate_size": 4096, "group_size": 128, "num_tokens": 32} + ) + + +@pytest.fixture(autouse=True) +def reset_config_manager_singleton(): + ConfigManager.reset_instance() + ConfigManager() + yield + ConfigManager.reset_instance() + + +class TestSiluAndMulPerBlockQuantCorrectness: + @pytest.mark.parametrize("num_tokens", [1, 7, 4096]) + @pytest.mark.parametrize("hidden_size", [1024, 2048, 5120]) + @pytest.mark.parametrize("group_size", [64, 128]) + @pytest.mark.parametrize("is_scale_transposed", [False, True]) + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) + @pytest.mark.parametrize("quant_dtype", [current_platform.fp8_dtype(), torch.int8]) + @pytest.mark.parametrize("has_scale_ub", [True, False]) + @pytest.mark.parametrize("seed", [0]) + def test_silu_and_mul_per_block_quant( + self, + num_tokens: int, + hidden_size: int, + group_size: int, + is_scale_transposed: bool, + dtype: torch.dtype, + quant_dtype: torch.dtype, + has_scale_ub: bool, + seed: int, + ) -> None: + skip_if_platform_unsupported("silu_and_mul_per_block_quant") + set_random_seed(seed) + + if hidden_size % group_size != 0: + return + + if has_scale_ub and quant_dtype != FP8_DTYPE: + # skip + return + + scale = 1 / hidden_size + x = torch.randn(num_tokens, 2 * hidden_size, dtype=dtype, device="cuda") * scale + + if has_scale_ub: + act = torch.nn.functional.silu(x[:, :hidden_size]) * x[:, hidden_size:] + act_abs = act.abs().float() + scale_ub = 0.5 * (act_abs.mean() + act_abs.amax()) + else: + scale_ub = None + + ref_out = torch.empty(num_tokens, hidden_size, device="cuda", dtype=quant_dtype) + + if is_scale_transposed: + ref_scales = torch.empty( + (hidden_size // group_size, x.shape[0]), + device="cuda", + dtype=torch.float32, + ).t() + else: + ref_scales = torch.empty( + (x.shape[0], hidden_size // group_size), + device="cuda", + dtype=torch.float32, + ) + + ops_out = ref_out.clone() + ops_scales = ref_scales.clone() + + baseline(ref_out, x, ref_scales, group_size, scale_ub, is_scale_transposed) + silu_and_mul_per_block_quant( + ops_out, x, ops_scales, group_size, scale_ub, is_scale_transposed + ) + + torch.testing.assert_close(ref_scales, ops_scales) + # allow 1 ULP difference + assert ( + ref_out.view(torch.uint8).to(torch.int16) + - ops_out.view(torch.uint8).to(torch.int16) + ).abs().max() <= 1 + + +class TestSiluAndMulPerBlockQuantIntegration: + def test_kernel_registration_integration(self): + from vllm.kernels.helion.register import get_registered_kernels + + registered_kernels = get_registered_kernels() + assert "silu_and_mul_per_block_quant" in registered_kernels + + kernel_wrapper = registered_kernels["silu_and_mul_per_block_quant"] + assert kernel_wrapper.op_name == "silu_and_mul_per_block_quant" + assert kernel_wrapper._config_picker is not None + assert kernel_wrapper._mutates_args == ["out", "scales"] + + def test_fake_impl_functionality(self): + skip_if_platform_unsupported("silu_and_mul_per_block_quant") + from vllm.kernels.helion.register import get_registered_kernels + + registered_kernels = get_registered_kernels() + kernel_wrapper = registered_kernels["silu_and_mul_per_block_quant"] + fake_impl = kernel_wrapper._fake_impl + + args = _generate_fake_input(16, 4096, 128) + assert fake_impl(*args) is None diff --git a/tests/kernels/mamba/test_cpu_short_conv.py b/tests/kernels/mamba/test_cpu_short_conv.py new file mode 100644 index 00000000000..cd32e0901a7 --- /dev/null +++ b/tests/kernels/mamba/test_cpu_short_conv.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from vllm.config import CompilationConfig, VllmConfig +from vllm.forward_context import set_forward_context +from vllm.model_executor.layers.mamba.short_conv import ShortConv +from vllm.model_executor.layers.utils import dispatch_cpu_unquantized_gemm +from vllm.platforms import current_platform +from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionMetadata + +if not current_platform.is_cpu(): + pytest.skip("skipping CPU-only tests", allow_module_level=True) + + +@pytest.fixture(autouse=True) +def mock_dist(): + with ( + patch( + "vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", + return_value=0, + ), + patch( + "vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", + return_value=1, + ), + patch( + "vllm.distributed.parallel_state.model_parallel_is_initialized", + return_value=True, + ), + patch( + "vllm.distributed.parallel_state.get_tp_group", + return_value=MagicMock(rank_in_group=0), + ), + ): + yield + + +@pytest.fixture +def vllm_config(): + # ShortConv only needs compilation_config from the current vLLM config, so a + # minimal config (model_config=None) avoids mocking ModelConfig and the + # associated VllmConfig validation churn. + return VllmConfig(compilation_config=CompilationConfig()) + + +def test_short_conv_forward_native_prefill(vllm_config): + prefix = "test_layer" + config = SimpleNamespace(conv_L_cache=4, conv_bias=True) + dim = 16 + + from vllm.config import set_current_vllm_config + + with set_current_vllm_config(vllm_config): + layer = ShortConv(config=config, dim=dim, layer_idx=0, prefix=prefix) + + layer.to("cpu") + dispatch_cpu_unquantized_gemm(layer.in_proj, remove_weight=False) + dispatch_cpu_unquantized_gemm(layer.out_proj, remove_weight=False) + + # Mock AttentionMetadata + num_prefills = 1 + num_prefill_tokens = 5 + query_start_loc_p = torch.tensor([0, 5], dtype=torch.int32) + state_indices_tensor_p = torch.tensor([0], dtype=torch.int32) + + # ShortConvAttentionMetadata + attn_metadata = ShortConvAttentionMetadata( + num_prefills=num_prefills, + num_prefill_tokens=num_prefill_tokens, + num_decodes=0, + num_decode_tokens=0, + num_reqs=1, + query_start_loc_p=query_start_loc_p, + has_initial_states_p=torch.tensor([False]), + state_indices_tensor_p=state_indices_tensor_p, + state_indices_tensor_d=torch.empty((0, 1), dtype=torch.int32), + num_accepted_tokens=None, + query_start_loc_d=None, + block_idx_last_scheduled_token=None, + block_idx_first_scheduled_token_p=None, + block_idx_last_computed_token=None, + block_idx_last_scheduled_token_prev_step=None, + num_computed_tokens_p=None, + seq_lens=torch.tensor([5]), + ) + + # Mock KV cache + # conv_state shape (num_blocks, L_cache - 1, dim) + conv_state = torch.zeros((1, config.conv_L_cache - 1, dim)) + layer.kv_cache = (conv_state,) + + hidden_states = torch.randn((num_prefill_tokens, dim)) + output = torch.zeros_like(hidden_states) + + attn_metadata_dict = {prefix: attn_metadata} + with set_forward_context(attn_metadata=attn_metadata_dict, vllm_config=vllm_config): + layer.forward_native(hidden_states, output) + + # Check if KV cache was updated + assert not torch.allclose(conv_state, torch.zeros_like(conv_state)) + + +def test_short_conv_forward_native_decode(vllm_config): + prefix = "test_layer_decode" + config = SimpleNamespace(conv_L_cache=4, conv_bias=True) + dim = 16 + + from vllm.config import set_current_vllm_config + + with set_current_vllm_config(vllm_config): + layer = ShortConv(config=config, dim=dim, layer_idx=0, prefix=prefix) + + layer.to("cpu") + dispatch_cpu_unquantized_gemm(layer.in_proj, remove_weight=False) + dispatch_cpu_unquantized_gemm(layer.out_proj, remove_weight=False) + + # Mock AttentionMetadata for 2 decode requests + num_decodes = 2 + state_indices_tensor_d = torch.tensor([0, 1], dtype=torch.int32) + + attn_metadata = ShortConvAttentionMetadata( + num_prefills=0, + num_prefill_tokens=0, + num_decodes=num_decodes, + num_decode_tokens=num_decodes, + num_reqs=num_decodes, + query_start_loc_p=None, + has_initial_states_p=None, + state_indices_tensor_p=torch.empty((0,), dtype=torch.int32), + state_indices_tensor_d=state_indices_tensor_d, + num_accepted_tokens=None, + query_start_loc_d=torch.tensor([0, 1, 2], dtype=torch.int32), + block_idx_last_scheduled_token=None, + block_idx_first_scheduled_token_p=None, + block_idx_last_computed_token=None, + block_idx_last_scheduled_token_prev_step=None, + num_computed_tokens_p=None, + seq_lens=torch.tensor([1, 1]), + ) + + # Mock KV cache (2 blocks for 2 requests) + conv_state = torch.randn((2, config.conv_L_cache - 1, dim)) + layer.kv_cache = (conv_state,) + + hidden_states = torch.randn((num_decodes, dim)) + output = torch.zeros_like(hidden_states) + + old_conv_state = conv_state.clone() + + attn_metadata_dict = {prefix: attn_metadata} + with set_forward_context(attn_metadata=attn_metadata_dict, vllm_config=vllm_config): + layer.forward_native(hidden_states, output) + + # Check if KV cache was updated + assert not torch.allclose(conv_state, old_conv_state) + + +def test_dispatch_cpu_unquantized_gemm_conv_layer(): + # Convolution layers have >2D weights; dispatch should skip them gracefully. + # Shape/dtype are AMX-pack safe (bf16, width==4, dim % block_size == 0) so + # the AMX prepack branch does not raise on AMX-capable CPUs. + class MockConvLayer(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter( + torch.randn(32, 1, 4, dtype=torch.bfloat16) + ) + self.bias = torch.nn.Parameter(torch.randn(32, dtype=torch.bfloat16)) + + layer = MockConvLayer() + # The ndim != 2 guard returns early without raising. + dispatch_cpu_unquantized_gemm(layer, remove_weight=False) + # No cpu_linear set — conv layers are handled elsewhere. + assert not hasattr(layer, "cpu_linear") diff --git a/tests/kernels/moe/test_cpu_int4_moe.py b/tests/kernels/moe/test_cpu_int4_moe.py index 05694eb08b2..04931e386f1 100644 --- a/tests/kernels/moe/test_cpu_int4_moe.py +++ b/tests/kernels/moe/test_cpu_int4_moe.py @@ -8,19 +8,21 @@ import pytest import torch import torch.nn.functional as F -from vllm.platforms import current_platform +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.experts.cpu_int4_moe import ( + CPUExpertsInt4, +) +from vllm.model_executor.layers.fused_moe.oracle.w4a8_int8 import ( + convert_to_w4a8_int8_moe_format, +) +from vllm.platforms import CpuArchEnum, current_platform from vllm.utils.torch_utils import set_random_seed -if not current_platform.is_cpu(): - pytest.skip("skipping CPU-only tests", allow_module_level=True) - -# Check if the dynamic_4bit_int_moe op is available -if not hasattr(torch.ops._C, "dynamic_4bit_int_moe"): - pytest.skip("dynamic_4bit_int_moe op not available", allow_module_level=True) - -# Check if KleidiAI ops are available -if not hasattr(torch.ops.aten, "_dyn_quant_pack_4bit_weight"): - pytest.skip("KleidiAI 4-bit ops not available", allow_module_level=True) +if ( + not current_platform.is_cpu() + or current_platform.get_cpu_architecture() != CpuArchEnum.ARM +): + pytest.skip("skipping Arm CPU-only tests", allow_module_level=True) # Tolerance for INT4 W4A8 @@ -34,49 +36,6 @@ def _silu_and_mul(x: torch.Tensor) -> torch.Tensor: return F.silu(x[..., :d]) * x[..., d:] -def _pack_int4_weight_to_kleidi( - int4_as_int8: torch.Tensor, - scales: torch.Tensor, - bias: torch.Tensor | None, - group_size: int, - in_features: int, - out_features: int, -) -> torch.Tensor: - """Pack INT4 weights (stored as int8 in [-8,7]) to KleidiAI format. - - Args: - int4_as_int8: [out, in] int8 tensor with values in [-8, 7] - scales: [out, in//group_size] or [out, 1] for channel-wise - bias: [out] optional bias - group_size: Quantization group size (-1 for channel-wise) - in_features: Input dimension - out_features: Output dimension - - Returns: - Packed weight tensor in KleidiAI format - """ - # Shift to unsigned nibble [0, 15] - tmp = int4_as_int8.add(8) - # Pack pairs along input dimension - uint8_nibbles = ((tmp[:, 1::2] << 4) | tmp[:, ::2]).to(torch.uint8) - - # Determine scale dtype based on group_size - scale_dtype = torch.float32 if group_size == -1 else torch.bfloat16 - scales_typed = scales.to(scale_dtype) - bias_typed = None if bias is None else bias.to(torch.float32) - - # Pack using KleidiAI op - actual_group_size = in_features if group_size == -1 else group_size - return torch.ops.aten._dyn_quant_pack_4bit_weight( - uint8_nibbles, - scales_typed, - bias_typed, - actual_group_size, - in_features, - out_features, - ) - - def _make_int4_moe_weights( E: int, N: int, @@ -124,59 +83,29 @@ def _make_int4_moe_weights( w13_bias = torch.randn(E, 2 * N, dtype=torch.float32) * 0.01 w2_bias = torch.randn(E, K, dtype=torch.float32) * 0.01 - # Pack weights for each expert - w13_packed_list = [] - w2_packed_list = [] + w13_packed, w2_packed, *_ = convert_to_w4a8_int8_moe_format( + w13_weight=w13_int4, + w2_weight=w2_int4, + w13_weight_scale=w13_scales, + w2_weight_scale=w2_scales, + group_size=group_size, + w13_bias=w13_bias if has_bias else None, + w2_bias=w2_bias if has_bias else None, + ) - for e in range(E): - w13_packed_list.append( - _pack_int4_weight_to_kleidi( - w13_int4[e], - w13_scales[e], - w13_bias[e] if (has_bias and w13_bias is not None) else None, - group_size, - K, - 2 * N, - ) - ) - w2_packed_list.append( - _pack_int4_weight_to_kleidi( - w2_int4[e], - w2_scales[e], - w2_bias[e] if (has_bias and w2_bias is not None) else None, - group_size, - N, - K, - ) - ) + if group_size == -1: + w13_scale = w13_scales.float() + w2_scale = w2_scales.float() + else: + w13_scale = w13_scales.float().repeat_interleave(group_size, dim=-1) + w2_scale = w2_scales.float().repeat_interleave(group_size, dim=-1) - w13_packed = torch.stack(w13_packed_list, dim=0) - w2_packed = torch.stack(w2_packed_list, dim=0) - - # Create reference dequantized weights - w13_ref = torch.zeros(E, 2 * N, K, dtype=torch.float32) - w2_ref = torch.zeros(E, K, N, dtype=torch.float32) - - for e in range(E): - # Dequantize w13 - for i in range(2 * N): - for j in range(K): - group_idx = 0 if group_size == -1 else (j // group_size) - w13_ref[e, i, j] = ( - w13_int4[e, i, j].float() * w13_scales[e, i, group_idx].float() - ) - if has_bias and w13_bias is not None: - w13_ref[e, i, j] += w13_bias[e, i].float() - - # Dequantize w2 - for i in range(K): - for j in range(N): - group_idx = 0 if group_size == -1 else (j // group_size) - w2_ref[e, i, j] = ( - w2_int4[e, i, j].float() * w2_scales[e, i, group_idx].float() - ) - if has_bias and w2_bias is not None: - w2_ref[e, i, j] += w2_bias[e, i].float() + w13_ref = w13_int4.float() * w13_scale + w2_ref = w2_int4.float() * w2_scale + if has_bias and w13_bias is not None: + w13_ref = w13_ref + w13_bias.float().unsqueeze(-1) + if has_bias and w2_bias is not None: + w2_ref = w2_ref + w2_bias.float().unsqueeze(-1) return w13_packed, w2_packed, w13_ref, w2_ref, w13_bias, w2_bias @@ -233,17 +162,20 @@ MoE_CONFIGS = [ (768, 2048, 16, 4, 64), ] SEEDS = [0, 42] +ACTIVATION_DTYPES = [torch.float32, torch.bfloat16, torch.float16] @pytest.mark.parametrize("M", NUM_TOKENS) @pytest.mark.parametrize("N,K,E,topk,group_size", MoE_CONFIGS) @pytest.mark.parametrize("seed", SEEDS) -def test_cpu_int4_moe_kernel(M, N, K, E, topk, group_size, seed): +@pytest.mark.parametrize("activation_dtype", ACTIVATION_DTYPES) +def test_cpu_int4_moe_kernel(M, N, K, E, topk, group_size, seed, activation_dtype): """Test dynamic_4bit_int_moe kernel against dequantized torch reference.""" set_random_seed(seed) + activation = MoEActivation.SILU # Generate input activations - a = torch.randn(M, K, dtype=torch.bfloat16) / (K**0.5) + a = torch.randn(M, K, dtype=activation_dtype) / (K**0.5) # Generate INT4 weights w13_packed, w2_packed, w13_ref, w2_ref, w13_bias, w2_bias = _make_int4_moe_weights( @@ -266,8 +198,6 @@ def test_cpu_int4_moe_kernel(M, N, K, E, topk, group_size, seed): ) # Test dynamic_4bit_int_moe kernel - # Activation kind: 1 = SwiGLU_Ug (SiLU(u)*g) for OAI-style - activation_kind = 1 apply_router_weight_on_input = False out = torch.ops._C.dynamic_4bit_int_moe( @@ -278,14 +208,14 @@ def test_cpu_int4_moe_kernel(M, N, K, E, topk, group_size, seed): w2_packed, K, # H (hidden_size / w2_out_features) N, # I (intermediate_size / w2_in_features) - 2 * N, # I2 (2*intermediate_size / w13_out_features) group_size, apply_router_weight_on_input, - activation_kind, + CPUExpertsInt4._activation_kind(activation), ) + assert out.dtype == activation_dtype torch.testing.assert_close( - ref_out.bfloat16(), + ref_out, out, atol=INT4_W4A8_ATOL, rtol=INT4_W4A8_RTOL, diff --git a/tests/kernels/moe/test_flashinfer_b12x_moe.py b/tests/kernels/moe/test_flashinfer_b12x_moe.py index 5aac3784ba4..b15cbcdd812 100644 --- a/tests/kernels/moe/test_flashinfer_b12x_moe.py +++ b/tests/kernels/moe/test_flashinfer_b12x_moe.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import pytest import torch @@ -8,8 +10,7 @@ from vllm.platforms import current_platform if not current_platform.is_device_capability_family(120): pytest.skip( - reason="FlashInfer CuteDSL SM12x MoE requires SM120 " - "(RTX Pro 6000 / DGX Spark).", + reason="FlashInfer B12x MoE requires SM120 (RTX Pro 6000 / DGX Spark).", allow_module_level=True, ) @@ -18,8 +19,8 @@ from vllm.utils.flashinfer import has_flashinfer_b12x_moe if not has_flashinfer_b12x_moe(): pytest.skip( reason=( - "FlashInfer cute_dsl_fused_moe_nvfp4 / convert_sf_to_mma_layout " - "not available in installed FlashInfer (needs PRs #3051 and #3066)." + "FlashInfer B12xMoEWrapper not available in installed " + "FlashInfer (needs PR #3080)." ), allow_module_level=True, ) @@ -40,7 +41,6 @@ from vllm.model_executor.layers.fused_moe.config import nvfp4_moe_quant_config from vllm.model_executor.layers.fused_moe.experts.flashinfer_b12x_moe import ( FlashInferB12xExperts, ) -from vllm.utils.flashinfer import flashinfer_convert_sf_to_mma_layout from vllm.utils.torch_utils import set_random_seed # Dimensions chosen to satisfy FP4 alignment requirements (k multiple of 256, @@ -59,7 +59,7 @@ def _reorder_gate_up_to_up_gate( ) -> tuple[torch.Tensor, torch.Tensor]: """Swap gate and up-projection halves along dim=1 to [up, gate] order. - The SM12x kernel expects weights in [up (w3), gate (w1)] order while the + The B12x kernel expects weights in [up (w3), gate (w1)] order while the BF16 reference uses [gate (w1), up (w3)]. This replicates the reordering done at model-load time by ``prepare_nvfp4_moe_layer_for_fi_or_cutlass``. """ @@ -70,6 +70,22 @@ def _reorder_gate_up_to_up_gate( ) +def _process_b12x_weights( + experts: FlashInferB12xExperts, + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + w1_scale_2: torch.Tensor, + w2_scale_2: torch.Tensor, +) -> None: + layer = SimpleNamespace( + w13_weight_scale=w1_scale, + w13_weight_scale_2=w1_scale_2, + w2_weight_scale=w2_scale, + w2_weight_scale_2=w2_scale_2, + ) + experts.process_weights_after_loading(layer) + + @pytest.mark.parametrize("m,n,k", MNK_FACTORS) @pytest.mark.parametrize("e", [8, 16]) @pytest.mark.parametrize("topk", [1, 2, 4]) @@ -174,22 +190,12 @@ def test_flashinfer_b12x_moe( moe_config=moe_config, quant_config=quant_config, ) - # In production, process_weights_after_loading computes these after - # normalizing block scales. In the test the scales are already in final - # form (global_scale=1.0), so we compute the MMA layouts directly. - num_experts_w1, m1, k1_sf = w1_blockscale.shape - experts.w1_sf_mma = flashinfer_convert_sf_to_mma_layout( - w1_blockscale.reshape(num_experts_w1 * m1, k1_sf), - m=m1, - k=k1_sf * 16, - num_groups=num_experts_w1, - ) - num_experts_w2, m2, k2_sf = w2_blockscale.shape - experts.w2_sf_mma = flashinfer_convert_sf_to_mma_layout( - w2_blockscale.reshape(num_experts_w2 * m2, k2_sf), - m=m2, - k=k2_sf * 16, - num_groups=num_experts_w2, + _process_b12x_weights( + experts, + w1_blockscale, + w2_blockscale, + ones_e, + ones_e, ) kernel = mk.FusedMoEKernel( @@ -224,5 +230,134 @@ def test_flashinfer_b12x_moe( torch.testing.assert_close(sm12x_output, torch_output, atol=2e-1, rtol=2e-1) +@pytest.mark.parametrize("m,n,k", MNK_FACTORS) +@pytest.mark.parametrize("e", [8, 16]) +@pytest.mark.parametrize("topk", [1, 2, 4]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@torch.inference_mode() +def test_flashinfer_b12x_moe_relu2( + m: int, + n: int, + k: int, + e: int, + topk: int, + dtype: torch.dtype, + workspace_init, +): + """Test FlashInferB12xExperts with ReLU2 (non-gated) activation. + + ReLU2 is used by Nemotron-H style models. Unlike the gated SiLU + path, w1 has shape [E, N, K] (not [E, 2N, K]) and the activation + is relu(x)^2 without a gate/up split. + """ + set_random_seed(7) + with set_current_vllm_config( + VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) + ): + a = torch.randn((m, k), device="cuda", dtype=dtype) / 10 + + # Non-gated: w1 shape is (e, n, k), not (e, 2n, k). + w1_bf16 = torch.randn((e, n, k), device="cuda", dtype=dtype) / 15 + w2_bf16 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 15 + + gs = torch.ones(1, device="cuda", dtype=torch.float32) + sf_vec_size = 16 + + # W1: no gate/up reordering for non-gated. + w1_flat = w1_bf16.reshape(e * n, k) + w1_q_flat, w1_sf_flat = fp4_quantize( + w1_flat, + global_scale=gs, + sf_vec_size=sf_vec_size, + is_sf_swizzled_layout=True, + ) + w1_q = w1_q_flat.view(e, n, k // 2) + w1_blockscale = w1_sf_flat.view(e, n, w1_sf_flat.shape[1]) + + w2_flat = w2_bf16.reshape(e * k, n) + w2_q_flat, w2_sf_flat = fp4_quantize( + w2_flat, + global_scale=gs, + sf_vec_size=sf_vec_size, + is_sf_swizzled_layout=True, + ) + w2_q = w2_q_flat.view(e, k, n // 2) + w2_blockscale = w2_sf_flat.view(e, k, w2_sf_flat.shape[1]) + + ones_e = torch.ones(e, device="cuda", dtype=torch.float32) + + quant_config = nvfp4_moe_quant_config( + g1_alphas=ones_e, + g2_alphas=ones_e, + a1_gscale=ones_e, + a2_gscale=ones_e, + w1_scale=w1_blockscale, + w2_scale=w2_blockscale, + ) + + moe_config = make_dummy_moe_config( + num_experts=e, + experts_per_token=topk, + hidden_dim=k, + intermediate_size=n, + in_dtype=dtype, + activation=MoEActivation.RELU2_NO_MUL, + ) + + experts = FlashInferB12xExperts( + moe_config=moe_config, + quant_config=quant_config, + ) + _process_b12x_weights( + experts, + w1_blockscale, + w2_blockscale, + ones_e, + ones_e, + ) + + kernel = mk.FusedMoEKernel( + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=False, + ), + experts, + inplace=False, + ) + + score = torch.randn((m, e), device="cuda", dtype=dtype) + topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False) + + b12x_output = kernel.apply( + hidden_states=a, + w1=w1_q, + w2=w2_q, + topk_weights=topk_weights, + topk_ids=topk_ids, + global_num_experts=e, + activation=MoEActivation.RELU2_NO_MUL, + apply_router_weight_on_input=False, + expert_map=None, + ) + + torch_output = torch_moe( + a, + w1_bf16, + w2_bf16, + score, + topk, + activation=MoEActivation.RELU2_NO_MUL, + ) + + torch.testing.assert_close( + b12x_output, + torch_output, + atol=2e-1, + rtol=2e-1, + ) + + if __name__ == "__main__": test_flashinfer_b12x_moe(16, 128, 256, 8, 2, torch.bfloat16) diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index 3f3bcebd11e..5fdcb8682f5 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -55,6 +55,7 @@ def make_dummy_moe_config( intermediate_size: int = 1, in_dtype: torch.dtype = torch.bfloat16, max_num_tokens: int = 512, + activation: MoEActivation = MoEActivation.SILU, ) -> FusedMoEConfig: """ This is a dummy config for the mk constructor interface @@ -73,7 +74,7 @@ def make_dummy_moe_config( else num_experts, num_logical_experts=num_experts, moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), - activation=MoEActivation.SILU, + activation=activation, in_dtype=in_dtype, device="cuda", routing_method=RoutingMethodType.TopK, diff --git a/tests/kernels/quantization/test_block_fp8.py b/tests/kernels/quantization/test_block_fp8.py index 4cb638e47af..c5eaa2f9321 100644 --- a/tests/kernels/quantization/test_block_fp8.py +++ b/tests/kernels/quantization/test_block_fp8.py @@ -11,6 +11,7 @@ from tests.kernels.quant_utils import ( native_per_token_group_quant_fp8, native_w8a8_block_matmul, ) +from tests.kernels.utils import fp8_ulp_distance from vllm.config import VllmConfig from vllm.model_executor.kernels.linear.scaled_mm.cutlass import cutlass_scaled_mm from vllm.model_executor.layers.quantization.utils.fp8_utils import ( @@ -93,7 +94,24 @@ def test_per_token_group_quant_fp8( tma_aligned_scales=tma_aligned_scales, ) - assert torch.allclose(out.to(torch.float32), ref_out.to(torch.float32), rtol=0.15) + if current_platform.is_rocm(): + # On gfx950 the Triton and PyTorch FP8 kernels can round in opposite + # directions when an element lands at the midpoint between two adjacent + # e4m3fn values (1-ULP tie-breaking). Verify: (1) no element is more + # than 1 FP8 ULP away, and (2) fewer than 0.05% of elements have any + # mismatch. Observed worst case across all parameter combos: 0.049%, + # max ULP = 1. + ulp = fp8_ulp_distance(out, ref_out) + assert (ulp <= 1).all(), ( + f"FP8 mismatch > 1 ULP: {int((ulp > 1).sum())} elements" + ) + assert float((ulp > 0).float().mean()) < 5e-4, ( + f"Too many 1-ULP mismatches: {int((ulp > 0).sum())}/{ulp.numel()}" + ) + else: + assert torch.allclose( + out.to(torch.float32), ref_out.to(torch.float32), rtol=0.15 + ) assert torch.allclose(scale, ref_scale) if column_major_scales: diff --git a/tests/lora/test_punica_ops.py b/tests/lora/test_punica_ops.py index be878472620..f94b54d9fb1 100644 --- a/tests/lora/test_punica_ops.py +++ b/tests/lora/test_punica_ops.py @@ -5,7 +5,6 @@ from threading import Lock import pytest import torch -import vllm.lora.ops.torch_ops as torch_ops import vllm.lora.ops.triton_ops as triton_ops from vllm.lora.ops.triton_ops import LoRAKernelMeta from vllm.lora.ops.triton_ops.utils import _LORA_A_PTR_DICT, _LORA_B_PTR_DICT @@ -22,6 +21,59 @@ def reset_device(reset_default_device): pass +@pytest.fixture(autouse=True) +def cleanup_fixture(): + """Override conftest's cleanup_fixture— not needed for punica tests.""" + yield + + +@pytest.fixture(autouse=True) +def dynamo_reset(): + """Override conftest's dynamo_reset — not needed for punica tests.""" + yield + + +def _cpu_bgmv_shrink( + inputs, lora_weight, output, seq_len_tensor, lora_indices, scaling=1.0 +): + """Memory-efficient shrink reference: per-LoRA matmul loop on CPU. + output[mask] = scaling * inputs[mask] @ weight.T""" + exploded = torch.repeat_interleave(lora_indices, seq_len_tensor) + for lid in exploded.unique(): + if lid < 0: + continue + mask = exploded == lid + inp = inputs[mask].to(output.dtype) + w = lora_weight[lid].to(output.dtype) + output[mask] = scaling * (inp @ w.T) + + +def _cpu_bgmv_expand( + inputs, + lora_weight, + output, + seq_len_tensor, + lora_indices, + offset=0, + add_inputs=False, +): + """Memory-efficient expand reference: per-LoRA matmul loop on CPU. + output[mask, offset:offset+n] (+)= inputs[mask] @ weight.T""" + exploded = torch.repeat_interleave(lora_indices, seq_len_tensor) + for lid in exploded.unique(): + if lid < 0: + continue + mask = exploded == lid + inp = inputs[mask].to(output.dtype) + w = lora_weight[lid].to(output.dtype) + n = w.shape[0] + result = inp @ w.T + if add_inputs: + output[mask, offset : offset + n] += result + else: + output[mask, offset : offset + n] = result + + # Utility shrink and expand operations used as reference implementations. def sgmv_shrink_for_nslices( nslices: int, @@ -36,22 +88,21 @@ def sgmv_shrink_for_nslices( num_tokens: int, scaling: float, ): - """ - Wrapper around torch_ops.sgmv_shrink that handles any nslices. - """ + """CPU reference for sgmv_shrink using per-LoRA matmul loop.""" + inp_cpu = inputs_tensor.cpu() + seq_cpu = seq_len_tensor.cpu() + idx_cpu = prompt_lora_mapping.cpu() + out_cpu = out_tensor.cpu() for index in range(nslices): - torch_ops.sgmv_shrink( - inputs_tensor, - lora_weights_lst[index], - out_tensor[index], - b_seq_start_loc, - seq_len_tensor, - prompt_lora_mapping, - batches, - max_seq_length, - num_tokens, - scaling, + _cpu_bgmv_shrink( + inp_cpu, + lora_weights_lst[index].cpu(), + out_cpu[index], + seq_cpu, + idx_cpu, + scaling=scaling, ) + out_tensor.copy_(out_cpu) def sgmv_expand_for_nslices( @@ -68,42 +119,21 @@ def sgmv_expand_for_nslices( num_tokens: int, add_inputs: bool, ) -> None: - """ - Wrapper around torch_ops.sgmv_expand that handles any nslices. - """ - if nslices == 1: - # Verify the torch's sgmv_expand op - torch_ops.sgmv_expand( - inputs_tensor[0], - lora_weights_lst[0], - out_tensor, - b_seq_start_loc, - seq_len_tensor, - prompt_lora_mapping, - batches, - max_seq_length, - num_tokens, + """CPU reference for sgmv_expand using per-LoRA matmul loop.""" + seq_cpu = seq_len_tensor.cpu() + idx_cpu = prompt_lora_mapping.cpu() + out_cpu = out_tensor.cpu() + for index in range(nslices): + _cpu_bgmv_expand( + inputs_tensor[index].cpu(), + lora_weights_lst[index].cpu(), + out_cpu, + seq_cpu, + idx_cpu, + offset=hidden_size * index, add_inputs=add_inputs, ) - else: - slice_offset = 0 - for index in range(nslices): - lora_weights = lora_weights_lst[index] - torch_ops.sgmv_expand_slice( - inputs_tensor[index], - lora_weights, - out_tensor, - b_seq_start_loc, - seq_len_tensor, - prompt_lora_mapping, - batches, - max_seq_length, - num_tokens, - slice_offset, - hidden_size, - add_inputs=add_inputs, - ) - slice_offset += hidden_size + out_tensor.copy_(out_cpu) _dict_lock = Lock() diff --git a/tests/models/registry.py b/tests/models/registry.py index 5d075d5b395..b5861b4a107 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -938,6 +938,7 @@ _MULTIMODAL_EXAMPLE_MODELS = { "HunYuanVLForConditionalGeneration": _HfExamplesInfo( "tencent/HunyuanOCR", hf_overrides={"num_experts": 0}, + is_available_online=False, ), "Idefics3ForConditionalGeneration": _HfExamplesInfo( "HuggingFaceM4/Idefics3-8B-Llama3", @@ -1491,8 +1492,6 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = { "EagleMistralLarge3ForCausalLM": _HfExamplesInfo( "mistralai/Mistral-Large-3-675B-Instruct-2512", speculative_model="mistralai/Mistral-Large-3-675B-Instruct-2512-Eagle", - # TODO: revert once figuring out OOM in CI - is_available_online=False, ), "LlamaForCausalLMEagle3": _HfExamplesInfo( "Qwen/Qwen3-8B", @@ -1557,6 +1556,12 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = { use_original_num_layers=True, ), # [MTP] + "BailingMoeV25MTPModel": _HfExamplesInfo( + "inclusionAI/Ring-2.5-1T", + speculative_model="inclusionAI/Ring-2.5-1T", + trust_remote_code=True, + is_available_online=False, + ), "DeepSeekMTPModel": _HfExamplesInfo( "luccafong/deepseek_mtp_main_random", speculative_model="luccafong/deepseek_mtp_draft_random", diff --git a/tests/models/transformers/__init__.py b/tests/models/transformers/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/models/transformers/fusers/__init__.py b/tests/models/transformers/fusers/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/models/transformers/fusers/test_linear.py b/tests/models/transformers/fusers/test_linear.py new file mode 100644 index 00000000000..3ead599234f --- /dev/null +++ b/tests/models/transformers/fusers/test_linear.py @@ -0,0 +1,480 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the Transformers modeling backend's linear fusers.""" + +import inspect +from types import MethodType, SimpleNamespace + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from vllm.model_executor.models.transformers.fuser import get_fuser +from vllm.model_executor.models.transformers.fusers import GLUFuser, QKVFuser + + +class SiluAndMulStub(nn.Module): + """Stand-in for vLLM's `SiluAndMul` (no vLLM config required).""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + return F.silu(x[..., :d]) * x[..., d:] + + +class NoDownGLU(nn.Module): + """`act(gate(x)) * up(x)` with no output projection -> `down_name` is None.""" + + def __init__(self, hidden: int = 16, inter: int = 32, bias: bool = False): + super().__init__() + self.gate_proj = nn.Linear(hidden, inter, bias=bias) + self.up_proj = nn.Linear(hidden, inter, bias=bias) + self.act_fn = nn.SiLU() + + def forward(self, x): + return self.act_fn(self.gate_proj(x)) * self.up_proj(x) + + +class GLUMLP(NoDownGLU): + """`down(act(gate(x)) * up(x))` — the canonical HF GLU MLP.""" + + def __init__(self, hidden: int = 16, inter: int = 32, bias: bool = False): + super().__init__(hidden, inter, bias) + self.down_proj = nn.Linear(inter, hidden, bias=bias) + + def forward(self, x): + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + +class ReversedGLUMLP(GLUMLP): + """`up(x) * act(gate(x))` — operands swapped (multiply is commutative).""" + + def forward(self, x): + return self.down_proj(self.up_proj(x) * self.act_fn(self.gate_proj(x))) + + +class NotAnMLP(nn.Module): + """Two linears but no activation*linear multiply -> must not match.""" + + def __init__(self): + super().__init__() + self.fc1 = nn.Linear(8, 8) + self.fc2 = nn.Linear(8, 8) + + def forward(self, x): + return self.fc2(self.fc1(x)) + + +class NotAnActGLUMLP(GLUMLP): + """GLU-shaped, but the "activation" is not a known activation module.""" + + def __init__(self): + super().__init__() + self.act_fn = nn.Dropout() + + +class UntraceableMLP(GLUMLP): + """Data-dependent control flow *before* the GLU -> no match.""" + + def forward(self, x): + if x.sum() > 0: # noqa: SIM108 - intentionally untraceable + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return x + + +class UntraceableTailGLUMLP(GLUMLP): + """Data-dependent control flow *after* the GLU -> still fusable.""" + + def forward(self, x): + y = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + if y.sum() > torch.inf: # intentionally untraceable + y = y * 0 + return y + + +class FakeAttention(nn.Module): + """HF v5-style attention: shape unpacking, dead KV branch, kwargs interface.""" + + is_causal = True + + def __init__( + self, + hidden: int = 32, + head_dim: int = 8, + heads: int = 4, + kv_heads: int = 4, + bias: bool = False, + layer_idx: int = 0, + ): + super().__init__() + self.config = SimpleNamespace(_attn_implementation="vllm") + self.layer_idx = layer_idx + self.head_dim = head_dim + self.scaling = head_dim**-0.5 + self.q_proj = nn.Linear(hidden, heads * head_dim, bias=bias) + self.k_proj = nn.Linear(hidden, kv_heads * head_dim, bias=bias) + self.v_proj = nn.Linear(hidden, kv_heads * head_dim, bias=bias) + self.o_proj = nn.Linear(heads * head_dim, hidden, bias=bias) + + def forward( + self, hidden_states, attention_mask=None, past_key_values=None, **kwargs + ): + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS + + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + if past_key_values is not None: + k, v = past_key_values.update(k, v, self.layer_idx) + attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, None + ) + attn_output, attn_weights = attention_interface( + self, q, k, v, attention_mask, scaling=self.scaling, **kwargs + ) + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + return self.o_proj(attn_output), attn_weights + + +class ReversedFakeAttention(FakeAttention): + """Projections computed in (v, k, q) order — q must still be identified.""" + + def forward( + self, hidden_states, attention_mask=None, past_key_values=None, **kwargs + ): + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS + + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, None + ) + attn_output, _ = attention_interface( + self, q, k, v, attention_mask, scaling=self.scaling, **kwargs + ) + return self.o_proj(attn_output.reshape(*input_shape, -1)), None + + +class ExtraProjAttention(FakeAttention): + """A second non-qkv linear of a different width -> `o_proj` still found.""" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.sink_proj = nn.Linear(self.head_dim, self.head_dim, bias=False) + + +class QKNormAttention(FakeAttention): + """OLMoE-style: a full-dim norm applied to the whole q/k projection output.""" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.q_norm = nn.RMSNorm(self.q_proj.out_features) + self.k_norm = nn.RMSNorm(self.k_proj.out_features) + + def forward( + self, hidden_states, attention_mask=None, past_key_values=None, **kwargs + ): + q = self.q_norm(self.q_proj(hidden_states)) + k = self.k_norm(self.k_proj(hidden_states)) + v = self.v_proj(hidden_states) + return self.o_proj(q + k + v), None + + +class PerHeadQKNormAttention(FakeAttention): + """Qwen3-style: a per-head norm (`head_dim`) applied after the head reshape.""" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.q_norm = nn.RMSNorm(self.head_dim) + self.k_norm = nn.RMSNorm(self.head_dim) + + def forward( + self, hidden_states, attention_mask=None, past_key_values=None, **kwargs + ): + shape = (*hidden_states.shape[:-1], -1, self.head_dim) + q = self.q_norm(self.q_proj(hidden_states).view(shape)) + k = self.k_norm(self.k_proj(hidden_states).view(shape)) + v = self.v_proj(hidden_states).view(shape) + return self.o_proj((q + k + v).flatten(-2)), None + + +class FakeSelfAttn(nn.Module): + """Stand-in for the vLLM `Attention` looked up in `attention_instances`.""" + + def __init__(self): + super().__init__() + self.impl = SimpleNamespace(scale=None) + + def forward(self, q, k, v): + # MHA-shaped stub: any deterministic combination of q/k/v will do + return q + 2 * k + 3 * v + + +@pytest.fixture(autouse=True) +def _clear_fuser_cache(): + get_fuser.cache_clear() + yield + get_fuser.cache_clear() + + +def _apply_glu_fuser_with_stubs(module: nn.Module, fuser: GLUFuser): + """Apply a fuser using plain stand-ins (merged `nn.Linear` + silu AndMul).""" + gate = module.get_submodule(fuser.gate_name) + up = module.get_submodule(fuser.up_name) + merged = nn.Linear( + gate.in_features, + gate.out_features + up.out_features, + bias=gate.bias is not None, + ) + with torch.no_grad(): + merged.weight.copy_(torch.cat([gate.weight, up.weight], dim=0)) + if gate.bias is not None: + merged.bias.copy_(torch.cat([gate.bias, up.bias], dim=0)) + setattr(module, fuser.merged_name, merged) + setattr(module, fuser.act_name, SiluAndMulStub()) + delattr(module, fuser.gate_name) + delattr(module, fuser.up_name) + module.forward = MethodType(fuser.fused_forward, module) + return module + + +def _apply_qkv_fuser_with_stubs(module: nn.Module, fuser: QKVFuser): + """Apply a fuser using a plain merged `nn.Linear` (no TP sharding).""" + q, k, v = ( + module.get_submodule(name) + for name in (fuser.q_name, fuser.k_name, fuser.v_name) + ) + merged = nn.Linear( + q.in_features, + q.out_features + k.out_features + v.out_features, + bias=q.bias is not None, + ) + with torch.no_grad(): + merged.weight.copy_(torch.cat([q.weight, k.weight, v.weight], dim=0)) + if q.bias is not None: + merged.bias.copy_(torch.cat([q.bias, k.bias, v.bias], dim=0)) + merged.split_sizes = [q.out_features, k.out_features, v.out_features] + setattr(module, fuser.merged_name, merged) + for name in (fuser.q_name, fuser.k_name, fuser.v_name): + delattr(module, name) + module.forward = MethodType(fuser.fused_forward, module) + return module + + +@pytest.mark.parametrize("mlp_cls", [GLUMLP, ReversedGLUMLP]) +@pytest.mark.parametrize("bias", [False, True]) +def test_detects_and_rewrites_glu(mlp_cls, bias): + with torch.device("meta"): + meta = mlp_cls(bias=bias) + fuser = get_fuser(meta) + assert isinstance(fuser, GLUFuser) + assert ( + fuser.gate_name, + fuser.up_name, + fuser.act_name, + fuser.down_name, + ) == ("gate_proj", "up_proj", "act_fn", "down_proj") + + # The rewritten forward references the merged projection instead of the + # sources; the rest of the forward is untouched. + names = fuser.fused_forward.__code__.co_names + assert "gate_up_proj" in names and "act_fn" in names and "down_proj" in names + assert not {"gate_proj", "up_proj"} & set(names) + + # Numerics: the fused forward must match the original on a real instance. + real = mlp_cls(bias=bias) + for p in real.parameters(): + nn.init.normal_(p, std=0.05) + x = torch.randn(4, 16) + expected = real(x) + fused = _apply_glu_fuser_with_stubs(real, fuser) + + # Fusion is in place: the module keeps its class and other attributes + assert fused is real and type(fused) is mlp_cls + torch.testing.assert_close(fused(x), expected, atol=1e-5, rtol=1e-5) + + +def test_glu_identifies_down_projection(): + """The row projection consuming `act(gate(x)) * up(x)` is identified. + + It is forced to `RowParallelLinear` in `update_attrs` so its sharded input + matches the column-parallel merged gate/up; `None` when there is no such + projection to force (fusion of gate/up still applies).""" + with torch.device("meta"): + assert get_fuser(GLUMLP()).down_name == "down_proj" + assert get_fuser(ReversedGLUMLP()).down_name == "down_proj" + assert get_fuser(NoDownGLU()).down_name is None + + +@pytest.mark.parametrize("attn_cls", [FakeAttention, ReversedFakeAttention]) +@pytest.mark.parametrize("kv_heads", [4, 2]) +def test_detects_and_rewrites_qkv(attn_cls, kv_heads): + if attn_cls is ReversedFakeAttention and kv_heads == 4: + pytest.skip("MHA q/k/v assignment is order-based by design") + with torch.device("meta"): + meta = attn_cls(kv_heads=kv_heads) + fuser = get_fuser(meta) + assert isinstance(fuser, QKVFuser) + # q (sharded differently under TP) must be identified exactly; k/v may be + # swapped for non-canonical compute order, which is numerically consistent + # because the weight mapping and the split indices follow the same + # assignment. + assert fuser.q_name == "q_proj" + assert {fuser.k_name, fuser.v_name} == {"k_proj", "v_proj"} + assert fuser.o_name == "o_proj" + + # The projections are merged; everything else stays live Python with its + # original semantics (branches, kwargs, attribute reads) + code = fuser.fused_forward.__code__ + names = code.co_names + assert "qkv_proj" in names and "split_sizes" in names and "o_proj" in names + assert not {"q_proj", "k_proj", "v_proj"} & set(names) + if attn_cls is FakeAttention: + assert "update" in names # the cache branch survives + assert code.co_flags & inspect.CO_VARKEYWORDS # **kwargs survives + + # Numerics: the fused forward must match the original on a real instance, + # with a different layer_idx than the traced instance (kv_heads == heads so + # the q/k/v stub combination is shape-compatible). + real = attn_cls(kv_heads=4, layer_idx=3) + for p in real.parameters(): + nn.init.normal_(p, std=0.05) + x = torch.randn(1, 5, 32) + attention_instances = {3: FakeSelfAttn()} + expected, _ = real(x, attention_instances=attention_instances) + fused = _apply_qkv_fuser_with_stubs(real, fuser) + + # Fusion is in place: the module keeps its class and other attributes + assert fused is real and type(fused) is attn_cls + assert fused.layer_idx == 3 and fused.is_causal and fused.config is not None + out, _ = fused(x, attention_instances=attention_instances) + torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5) + + +def test_qkv_identifies_output_projection(): + with torch.device("meta"): + assert get_fuser(FakeAttention()).o_name == "o_proj" + assert get_fuser(ReversedFakeAttention()).o_name == "o_proj" + assert get_fuser(ExtraProjAttention()).o_name == "o_proj" + # Norm children (q_norm/k_norm) must not disturb o_proj identification. + assert get_fuser(QKNormAttention()).o_name == "o_proj" + assert get_fuser(PerHeadQKNormAttention()).o_name == "o_proj" + + +def test_fuser_is_cached_per_class(): + with torch.device("meta"): + fuser_a = get_fuser(GLUMLP()) + fuser_b = get_fuser(GLUMLP()) + assert fuser_a is fuser_b + assert GLUMLP in get_fuser.cache + + +@pytest.mark.parametrize("cls", [NotAnMLP, UntraceableMLP]) +def test_non_matching_modules_return_none(cls): + with torch.device("meta"): + module = cls() + assert get_fuser(module) is None + + +def test_untraceable_tail_still_fuses(): + with torch.device("meta"): + meta = UntraceableTailGLUMLP() + fuser = get_fuser(meta) + assert isinstance(fuser, GLUFuser) + + # Numerics: the live tail must survive the rewrite + real = UntraceableTailGLUMLP() + for p in real.parameters(): + nn.init.normal_(p, std=0.05) + x = torch.randn(4, 16) + expected = real(x) + fused = _apply_glu_fuser_with_stubs(real, fuser) + torch.testing.assert_close(fused(x), expected, atol=1e-5, rtol=1e-5) + + +def test_weight_mappings_are_scoped_to_fused_prefixes(): + from vllm.model_executor.models.utils import WeightsMapper + + with torch.device("meta"): + glu_fuser = get_fuser(GLUMLP()) + qkv_fuser = get_fuser(FakeAttention()) + + mapper = WeightsMapper() + for prefix in ("model.layers.0.mlp", "model.layers.1.mlp"): + mapper.orig_to_new_stacked.update(glu_fuser.orig_to_new_stacked(prefix)) + mapper.orig_to_new_stacked.update( + qkv_fuser.orig_to_new_stacked("model.layers.0.self_attn") + ) + + names = [ + "model.layers.0.mlp.gate_proj.weight", + "model.layers.0.mlp.up_proj.weight", + "model.layers.1.mlp.gate_proj.weight", + "model.layers.0.self_attn.q_proj.weight", + "model.layers.0.self_attn.k_proj.weight", + "model.layers.0.self_attn.v_proj.weight", + # Unfused modules at other prefixes must be left untouched. + "model.layers.2.mlp.experts.0.gate_proj.weight", + "model.layers.1.self_attn.q_proj.weight", + ] + # `apply` rewrites the name and stamps the shard id onto each tensor. + weights = [(name, torch.empty(0)) for name in names] + mapped = list(mapper.apply(weights)) + mapped_names = [name for name, _ in mapped] + shard_ids = [getattr(data, "shard_id", None) for _, data in mapped] + + assert mapped_names == [ + "model.layers.0.mlp.gate_up_proj.weight", + "model.layers.0.mlp.gate_up_proj.weight", + "model.layers.1.mlp.gate_up_proj.weight", + "model.layers.0.self_attn.qkv_proj.weight", + "model.layers.0.self_attn.qkv_proj.weight", + "model.layers.0.self_attn.qkv_proj.weight", + # Only the exact fused layers are remapped; everything else is untouched. + "model.layers.2.mlp.experts.0.gate_proj.weight", + "model.layers.1.self_attn.q_proj.weight", + ] + assert shard_ids == [0, 1, 0, "q", "k", "v", None, None] + + # The fused layers are exposed to the quantization machinery via their + # original constituent projection names (what the checkpoint stores). + assert glu_fuser.packed_modules_mapping == { + "gate_up_proj": ["gate_proj", "up_proj"], + } + assert qkv_fuser.packed_modules_mapping == { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + } + + +@pytest.mark.parametrize("cls", [NotAnMLP, NotAnActGLUMLP]) +def test_unfusable_modules_are_not_fused(cls, default_vllm_config): + with torch.device("meta"): + module = cls() + fuser = get_fuser(module) + # Either no pattern matches the class, or this instance fails validation + # (`recursive_replace` gates fusion and its weight mappings on `validate`) + model_config = default_vllm_config.model_config + assert fuser is None or not fuser.validate(module, model_config) + + +def test_act_and_mul_derived_from_module(default_vllm_config): + from transformers.activations import GELUTanh, SiLUActivation + + from vllm.model_executor.layers.activation import GeluAndMul, SiluAndMul + + assert isinstance(GLUFuser._get_act_and_mul(nn.SiLU()), SiluAndMul) + assert isinstance(GLUFuser._get_act_and_mul(SiLUActivation()), SiluAndMul) + gelu_tanh = GLUFuser._get_act_and_mul(GELUTanh()) + assert isinstance(gelu_tanh, GeluAndMul) and gelu_tanh.approximate == "tanh" + gelu = GLUFuser._get_act_and_mul(nn.GELU()) + assert isinstance(gelu, GeluAndMul) and gelu.approximate == "none" + # Not activations at all -> no fusion + assert GLUFuser._get_act_and_mul_name(nn.Dropout()) is None + assert GLUFuser._get_act_and_mul_name(nn.LayerNorm(8)) is None + with pytest.raises(ValueError, match="No AndMul equivalent"): + GLUFuser._get_act_and_mul(nn.Dropout()) diff --git a/tests/models/transformers/fusers/test_moe.py b/tests/models/transformers/fusers/test_moe.py new file mode 100644 index 00000000000..04eadac3f78 --- /dev/null +++ b/tests/models/transformers/fusers/test_moe.py @@ -0,0 +1,299 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the Transformers modeling backend's MoE fuser.""" + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from vllm.model_executor.models.transformers.fusers import MoEBlockFuser + +from .test_linear import GLUMLP + + +class TopKRouter(nn.Module): + """HF v5 top-k router: `linear -> softmax -> topk (-> renorm)`.""" + + def __init__(self, num_experts=8, hidden=16, top_k=2, sigmoid=False): + super().__init__() + self.top_k = top_k + self.sigmoid = sigmoid + self.weight = nn.Parameter(torch.zeros(num_experts, hidden)) + + def forward(self, hidden_states): + logits = F.linear(hidden_states, self.weight) + scores = torch.sigmoid(logits) if self.sigmoid else F.softmax(logits, dim=-1) + value, index = torch.topk(scores, self.top_k, dim=-1) + value = value / value.sum(dim=-1, keepdim=True) + return logits, value, index + + +class CorrectionRouter(nn.Module): + """Grouped router with a score-correction bias buffer (DeepSeek-V3) -> declined.""" + + def __init__(self, num_experts=8, hidden=16): + super().__init__() + self.weight = nn.Parameter(torch.zeros(num_experts, hidden)) + self.register_buffer("e_score_correction_bias", torch.zeros(num_experts)) + + def forward(self, hidden_states): + logits = F.linear(hidden_states, self.weight) + scores = torch.sigmoid(logits) + self.e_score_correction_bias + _, index = torch.topk(scores, 2, dim=-1) + return logits, scores, index + + +class BiasedRouter(TopKRouter): + """A valid top-k router but not `weight`-only (extra `bias` param) -> declined.""" + + def __init__(self): + super().__init__() + self.bias = nn.Parameter(torch.zeros(8)) + + def forward(self, hidden_states): + logits = F.linear(hidden_states, self.weight) + self.bias + scores = F.softmax(logits, dim=-1) + value, index = torch.topk(scores, self.top_k, dim=-1) + return logits, value, index + + +class DisconnectedRouter(TopKRouter): + """linear+softmax+top-k present but top-k ignores the logits -> not a router.""" + + def forward(self, hidden_states): + logits = F.linear(hidden_states, self.weight) + _ = F.softmax(logits, dim=-1) # scored, but not consumed by top-k + value, index = torch.topk(hidden_states, self.top_k, dim=-1) + return logits, value, index + + +class MoEExperts(nn.Module): + """Packed experts (3D weights); only its name (`experts`) matters here.""" + + def __init__(self, num_experts=8, hidden=16, inter=32): + super().__init__() + self.gate_up_proj = nn.Parameter(torch.zeros(num_experts, 2 * inter, hidden)) + self.down_proj = nn.Parameter(torch.zeros(num_experts, hidden, inter)) + + def forward(self, hidden_states, index, weights): + return hidden_states + + +class MoEBlock(nn.Module): + """Single-tensor MoE block (Qwen3-style); subclasses override `_shared`.""" + + def __init__(self, router_cls=TopKRouter): + super().__init__() + self.experts = MoEExperts() + self.gate = router_cls() + + def _shared(self, x, logits): + """The term added to the experts' output (none for a plain block).""" + return 0 + + def forward(self, hidden_states): + x = hidden_states.reshape(-1, hidden_states.shape[-1]) + logits, weights, index = self.gate(x) + out = self.experts(x, index, weights) + self._shared(x, logits) + return out.reshape(hidden_states.shape) + + +class MoEBlockNoShared(MoEBlock): + """No shared-expert child but a gate-derived add -> trace skipped, still fuses.""" + + def _shared(self, x, logits): + return logits.sum() + + +class MoEBlockShared(MoEBlock): + """A block with a shared expert and its sigmoid gate (Qwen2-style).""" + + def __init__(self): + super().__init__() + self.shared_expert = GLUMLP() + self.shared_expert_gate = nn.Linear(16, 1, bias=False) + + def _shared(self, x, logits): + return torch.sigmoid(self.shared_expert_gate(x)) * self.shared_expert(x) + + +class MoEBlockSharedNoGate(MoEBlock): + """A block with an ungated shared expert -> native, shared passed through.""" + + def __init__(self): + super().__init__() + self.shared_expert = GLUMLP() + + def _shared(self, x, logits): + return self.shared_expert(x) + + +class MoEBlockTuple(MoEBlock): + """A tuple-returning block (gpt-oss-style) -> must decline.""" + + def forward(self, hidden_states): + x = hidden_states.reshape(-1, hidden_states.shape[-1]) + _, weights, index = self.gate(x) + return self.experts(x, index, weights), index + + +class MoEBlockTupleVar(MoEBlock): + """Returns a name bound to a tuple, not a literal tuple -> must still decline.""" + + def forward(self, hidden_states): + x = hidden_states.reshape(-1, hidden_states.shape[-1]) + _, weights, index = self.gate(x) + result = self.experts(x, index, weights), index + return result + + +class MoEBlockNestedTupleReturn(MoEBlock): + """Tuple `return` in a nested helper; block returns one tensor -> still fuses.""" + + def forward(self, hidden_states): + def keep(a, b): + return a, b + + x = hidden_states.reshape(-1, hidden_states.shape[-1]) + _, weights, index = self.gate(x) + out, _ = keep(self.experts(x, index, weights), index) + return out.reshape(hidden_states.shape) + + +class PlainMLP(nn.Module): + """A non-GLU FFN: `down(act(up(x)))`, no gating multiply.""" + + def __init__(self, hidden: int = 16, inter: int = 32): + super().__init__() + self.up_proj = nn.Linear(hidden, inter, bias=False) + self.down_proj = nn.Linear(inter, hidden, bias=False) + self.act_fn = nn.SiLU() + + def forward(self, x): + return self.down_proj(self.act_fn(self.up_proj(x))) + + +class MoEBlockSharedNonGLU(MoEBlock): + """A non-GLU shared expert -> detected by dataflow (no gate/up merge).""" + + def __init__(self): + super().__init__() + self.shared_expert = PlainMLP() + + def _shared(self, x, logits): + return self.shared_expert(x) + + +class MoEBlockUnaccounted(MoEBlock): + """A weight-bearing child outside the fused dataflow (pre-router) -> declined.""" + + def __init__(self): + super().__init__() + self.extra = nn.Linear(16, 16, bias=False) + + def forward(self, hidden_states): + x = self.extra(hidden_states.reshape(-1, hidden_states.shape[-1])) + _, weights, index = self.gate(x) + return self.experts(x, index, weights).reshape(hidden_states.shape) + + +class BufferScale(nn.Module): + """A stateful child carrying only a buffer (no parameters).""" + + def __init__(self, hidden: int = 16): + super().__init__() + self.register_buffer("scale", torch.ones(hidden)) + + def forward(self, x): + return x * self.scale + + +class MoEBlockUnaccountedBuffer(MoEBlockUnaccounted): + """Like `MoEBlockUnaccounted`, but the extra child holds only a buffer.""" + + def __init__(self): + super().__init__() + self.extra = BufferScale() + + +@pytest.mark.parametrize("sigmoid", [False, True]) +def test_moe_fuser_detects_router(sigmoid): + with torch.device("meta"): + block = MoEBlock(lambda: TopKRouter(sigmoid=sigmoid)) + fuser = MoEBlockFuser.match(block, "experts") + assert isinstance(fuser, MoEBlockFuser) + assert fuser.gate_name == "gate" + assert fuser.scoring_func == ("sigmoid" if sigmoid else "softmax") + assert fuser.shared_name is None and fuser.shared_gate_name is None + + +def test_moe_fuser_detects_shared_experts(): + with torch.device("meta"): + block = MoEBlockShared() + fuser = MoEBlockFuser.match(block, "experts") + assert isinstance(fuser, MoEBlockFuser) + assert fuser.shared_name == "shared_expert" + assert fuser.shared_gate_name == "shared_expert_gate" + + +def test_moe_fuser_skips_shared_detection_without_extra_children(): + """With only experts and gate, shared-expert detection (and its block trace) + is skipped, so a gate-derived add is not misread as a shared expert.""" + with torch.device("meta"): + block = MoEBlockNoShared() + fuser = MoEBlockFuser.match(block, "experts") + assert isinstance(fuser, MoEBlockFuser) + assert fuser.shared_name is None and fuser.shared_gate_name is None + + +def test_moe_fuser_shared_without_gate(): + with torch.device("meta"): + block = MoEBlockSharedNoGate() + fuser = MoEBlockFuser.match(block, "experts") + assert isinstance(fuser, MoEBlockFuser) + assert fuser.shared_name == "shared_expert" + assert fuser.shared_gate_name is None + + +def test_moe_fuser_detects_non_glu_shared_expert(): + with torch.device("meta"): + block = MoEBlockSharedNonGLU() + fuser = MoEBlockFuser.match(block, "experts") + assert isinstance(fuser, MoEBlockFuser) + # Recognised by dataflow (added to the experts' output), though not a GLU. + assert fuser.shared_name == "shared_expert" + assert fuser.shared_gate_name is None + + +@pytest.mark.parametrize( + "block_cls", + [ + lambda: MoEBlock(CorrectionRouter), # score-correction buffer (grouped) + lambda: MoEBlock(BiasedRouter), # router not weight-only (extra param) + MoEBlockTuple, # tuple-returning block (e.g. gpt-oss) + MoEBlockTupleVar, # tuple returned via a name binding, not a literal + MoEBlockUnaccounted, # weight-bearing child outside the fused dataflow + MoEBlockUnaccountedBuffer, # buffer-only child outside the fused dataflow + ], +) +def test_moe_fuser_declines_unsupported(block_cls): + with torch.device("meta"): + block = block_cls() + assert MoEBlockFuser.match(block, "experts") is None + + +def test_moe_fuser_ignores_nested_returns(): + """A tuple `return` inside a nested helper must not decline a block whose own + forward returns a single tensor.""" + with torch.device("meta"): + block = MoEBlockNestedTupleReturn() + assert isinstance(MoEBlockFuser.match(block, "experts"), MoEBlockFuser) + + +def test_moe_fuser_router_requires_connected_dataflow(): + """A gate with linear + softmax + top-k present but not wired as a router + (top-k selects over the input, not the scored logits) is not detected.""" + with torch.device("meta"): + block = MoEBlock(DisconnectedRouter) + assert MoEBlockFuser.match(block, "experts") is None diff --git a/tests/models/transformers/fusers/test_rms_norm.py b/tests/models/transformers/fusers/test_rms_norm.py new file mode 100644 index 00000000000..6497b98c4ba --- /dev/null +++ b/tests/models/transformers/fusers/test_rms_norm.py @@ -0,0 +1,227 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the Transformers modeling backend's RMSNorm fuser.""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from vllm.model_executor.models.transformers.fuser import get_fuser +from vllm.model_executor.models.transformers.fusers import RMSNormFuser + + +class RMSNorm(nn.Module): + """The canonical HF RMSNorm: `weight * x * rsqrt(mean(x**2) + eps)`.""" + + def __init__(self, hidden: int = 16, eps: float = 1e-5, weight: bool = True): + super().__init__() + if weight: + self.weight = nn.Parameter(torch.ones(hidden)) + self.variance_epsilon = eps + + def _rms(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon) + + def forward(self, x): + return self.weight * self._rms(x.to(torch.float32)).to(x.dtype) + + +class GemmaRMSNorm(RMSNorm): + """Zero-centered weight: `(1 + weight) * normalized`.""" + + def __init__(self, hidden: int = 16, eps: float = 1e-6): + super().__init__(hidden, eps) + self.weight = nn.Parameter(torch.zeros(hidden)) + + def forward(self, x): + return (1.0 + self.weight) * self._rms(x.to(torch.float32)).to(x.dtype) + + +class WeightlessRMSNorm(RMSNorm): + """No scale parameter (e.g. Gemma3n `with_scale=False`).""" + + def __init__(self, hidden: int = 16, eps: float = 1e-6): + super().__init__(hidden, eps, weight=False) + + def forward(self, x): + return self._rms(x.to(torch.float32)).to(x.dtype) + + +class LayerNorm(RMSNorm): + """An RMSNorm not named `*RMSNorm`, keeping the input dtype (no upcast).""" + + def __init__(self, hidden: int = 16, eps: float = 1e-6): + super().__init__(hidden, eps) + + def forward(self, x): + return self.weight * self._rms(x) + + +class NotAnRMSNorm(RMSNorm): + """Mean-subtracting LayerNorm-like math -> not an RMSNorm.""" + + def __init__(self, hidden: int = 16, eps: float = 1e-6): + super().__init__(hidden, eps) + + def forward(self, x): + x = x - x.mean(-1, keepdim=True) + variance = x.var(-1, keepdim=True) + return self.weight * x / torch.sqrt(variance + self.variance_epsilon) + + +class GatedRMSNorm(RMSNorm): + """Second input and tail compute -> not an RMSNorm.""" + + def forward(self, x, gate=None): + normed = self.weight * self._rms(x.to(torch.float32)).to(x.dtype) + return normed * F.silu(gate) + + +class GatedFusedRMSNorm(nn.Module): + """Same as GatedRMSNorm, but built on the fused `rms_norm` op -> not an RMSNorm.""" + + def __init__(self, hidden: int = 16, eps: float = 1e-5): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden)) + self.eps = eps + + def forward(self, x, gate=None): + return F.rms_norm(x, (x.shape[-1],), self.weight, self.eps) * F.silu(gate) + + +class UntraceableGatedRMSNorm(RMSNorm): + """Tracer can't see tail compute in forward, but still has a second input (gate).""" + + def forward(self, x, gate=None): + normed = self.weight * self._rms(x.to(torch.float32)).to(x.dtype) + if gate.sum() > 0: # untraceable -> partial graph, no visible tail + normed = normed * F.silu(gate) + return normed + + +@pytest.mark.parametrize( + "cls,eps,zero_centered", + [ + (RMSNorm, 1e-5, False), + (GemmaRMSNorm, 1e-6, True), + (WeightlessRMSNorm, 1e-6, False), + (LayerNorm, 1e-6, False), + (torch.nn.RMSNorm, 1e-5, False), # fused `F.rms_norm` op + ], +) +def test_detects_rms_norm_variants(cls, eps, zero_centered): + with torch.device("meta"): + fuser = get_fuser(cls(16, eps=eps)) + assert isinstance(fuser, RMSNormFuser) + assert fuser.zero_centered == zero_centered + + +@pytest.mark.parametrize("cls", [NotAnRMSNorm, nn.LayerNorm, nn.SiLU]) +def test_non_rms_norms_are_not_matched(cls): + with torch.device("meta"): + module = cls(16) if cls is nn.LayerNorm else cls() + assert not isinstance(get_fuser(module), RMSNormFuser) + + +@pytest.mark.parametrize( + "cls", [GatedRMSNorm, GatedFusedRMSNorm, UntraceableGatedRMSNorm] +) +def test_gated_rms_norm_is_not_fused(cls): + with torch.device("meta"): + assert not isinstance(get_fuser(cls()), RMSNormFuser) + + +@pytest.mark.parametrize( + "cls,expected,zero_centered", + [ + (RMSNorm, "RMSNorm", False), + (GemmaRMSNorm, "GemmaRMSNorm", True), + (WeightlessRMSNorm, "RMSNorm", False), + ], +) +def test_rms_norm_builds_vllm_class(cls, expected, zero_centered, default_vllm_config): + from vllm.model_executor.layers.layernorm import GemmaRMSNorm as VLLMGemmaRMSNorm + from vllm.model_executor.layers.layernorm import RMSNorm as VLLMRMSNorm + + # `default_vllm_config` supplies the config context the CustomOp needs; the + # weightless path reads hidden size from the model config, so stub it. + model_config = SimpleNamespace(get_hidden_size=lambda: 16) + with torch.device("meta"): + module = cls() + fuser = get_fuser(module) + built = fuser.fuse(module, "norm", model_config, None) + from vllm.model_executor.models.transformers.fusers.rms_norm import ( + TPAwareNormMixin, + ) + + types_by_name = {"RMSNorm": VLLMRMSNorm, "GemmaRMSNorm": VLLMGemmaRMSNorm} + assert isinstance(built, types_by_name[expected]) + assert isinstance(built, TPAwareNormMixin) # fused norms self-correct under TP + assert built.variance_epsilon == module.variance_epsilon + assert isinstance(built.weight, nn.Parameter) == ( + getattr(module, "weight", None) is not None + ) + + +def test_fused_rms_norm_op_default_eps(default_vllm_config): + """`torch.nn.RMSNorm` (a single `F.rms_norm` call) matches via the fast path; + its default `eps=None` resolves to `finfo(dtype).eps` in `fuse`.""" + from vllm.model_executor.layers.layernorm import RMSNorm as VLLMRMSNorm + + with torch.device("meta"): + module = torch.nn.RMSNorm(16) # forward is a single `F.rms_norm` call + fuser = get_fuser(module) + assert isinstance(fuser, RMSNormFuser) + assert not fuser.zero_centered + model_config = SimpleNamespace(get_hidden_size=lambda: 16, dtype=torch.float32) + built = fuser.fuse(module, "norm", model_config, None) + assert isinstance(built, VLLMRMSNorm) + assert built.variance_epsilon == torch.finfo(torch.float32).eps + + +def test_eps_is_derived_per_instance(default_vllm_config): + """Two instances of the same norm class with different eps must fuse to their + own eps: the type-cached fuser holds only structure, not this value.""" + model_config = SimpleNamespace(get_hidden_size=lambda: 16) + with torch.device("meta"): + for eps in (1e-5, 1e-6): + module = RMSNorm(16, eps=eps) + built = get_fuser(module).fuse(module, "norm", model_config, None) + assert built.variance_epsilon == eps + + +def test_fused_norm_is_gather_capable(default_vllm_config): + """Every fused norm is emitted gather-capable, so a norm on a head-sharded + projection (OLMoE-style) self-corrects at runtime with no QKV-specific + plumbing. A full-width input skips the gather and equals a plain norm.""" + from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm + from vllm.model_executor.models.transformers.fusers import rms_norm + + torch.manual_seed(0) + x = torch.randn(4, 16) + for gathered_cls, plain_cls in [ + (rms_norm.TPAwareRMSNorm, RMSNorm), + (rms_norm.TPAwareGemmaRMSNorm, GemmaRMSNorm), + ]: + gathered = gathered_cls(hidden_size=16, eps=1e-6) + assert isinstance(gathered, rms_norm.TPAwareNormMixin) + plain = plain_cls(hidden_size=16, eps=1e-6) + with torch.no_grad(): + weight = torch.randn(16) + gathered.weight.copy_(weight) + plain.weight.copy_(weight) + torch.testing.assert_close(gathered(x), plain(x)) + + +def test_gathered_norm_rejects_uneven_sharding(default_vllm_config): + """A sharded input (narrower than the full-width weight) that does not tile + the weight evenly across ranks is rejected before any collective.""" + from vllm.model_executor.models.transformers.fusers import rms_norm + + norm = rms_norm.TPAwareRMSNorm(hidden_size=8, eps=1e-6) + norm.tp_size = 2 # emulate TP=2 without a real process group + with pytest.raises(ValueError, match="does not tile it evenly"): + norm(torch.randn(2, 3)) # 3 * 2 != 8 diff --git a/tests/models/test_transformers.py b/tests/models/transformers/test_backend.py similarity index 82% rename from tests/models/test_transformers.py rename to tests/models/transformers/test_backend.py index eadc3534c37..a3eea1783f5 100644 --- a/tests/models/test_transformers.py +++ b/tests/models/transformers/test_backend.py @@ -6,10 +6,16 @@ from typing import Any import pytest -from ..conftest import HfRunner, VllmRunner -from ..utils import multi_gpu_test, prep_prompts -from .registry import HF_EXAMPLE_MODELS -from .utils import check_embeddings_close, check_logprobs_close +from ...conftest import HfRunner, VllmRunner +from ...utils import multi_gpu_test, prep_prompts +from ..registry import HF_EXAMPLE_MODELS +from ..utils import check_embeddings_close, check_logprobs_close + + +@pytest.fixture(scope="function", autouse=True) +def enable_pickle(monkeypatch): + """`LLM.apply_model` requires pickling a function.""" + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") def get_model(arch: str) -> str: @@ -18,6 +24,17 @@ def get_model(arch: str) -> str: return model_info.default +def get_num_fused(model) -> tuple[int, int]: + from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + ) + + glu = sum(isinstance(m, MergedColumnParallelLinear) for m in model.modules()) + qkv = sum(isinstance(m, QKVParallelLinear) for m in model.modules()) + return glu, qkv + + def check_implementation( runner_ref: type[HfRunner | VllmRunner], runner_test: type[VllmRunner], @@ -25,6 +42,7 @@ def check_implementation( model: str, kwargs_ref: dict[str, Any] | None = None, kwargs_test: dict[str, Any] | None = None, + num_fused: tuple[int, int] = (1, 1), **kwargs, ): if kwargs_ref is None: @@ -41,6 +59,12 @@ def check_implementation( model_config = model_test.llm.llm_engine.model_config assert model_config.using_transformers_backend() + num_layers = model_config.hf_config.get_text_config().num_hidden_layers + expected_glu, expected_qkv = num_fused + for num_glu, num_qkv in model_test.apply_model(get_num_fused): + assert num_glu == expected_glu * num_layers + assert num_qkv == expected_qkv * num_layers + outputs_test = model_test.generate_greedy_logprobs(*args) with runner_ref(model, **kwargs_ref) as model_ref: @@ -58,11 +82,11 @@ def check_implementation( @pytest.mark.parametrize( - "model,model_impl", + "model,model_impl,num_fused", [ - ("meta-llama/Llama-3.2-1B-Instruct", "transformers"), - ("hmellor/Ilama-3.2-1B", "auto"), # CUSTOM CODE - ("allenai/OLMoE-1B-7B-0924", "transformers"), # MoE + ("meta-llama/Llama-3.2-1B-Instruct", "transformers", (1, 1)), + ("hmellor/Ilama-3.2-1B", "auto", (1, 1)), # CUSTOM CODE + ("allenai/OLMoE-1B-7B-0924", "transformers", (0, 1)), # MoE ], ) # trust_remote_code=True by default def test_models( @@ -71,6 +95,7 @@ def test_models( example_prompts: list[str], model: str, model_impl: str, + num_fused: tuple[int, int], ) -> None: import transformers from packaging.version import Version @@ -84,7 +109,12 @@ def test_models( ) check_implementation( - hf_runner, vllm_runner, example_prompts, model, model_impl=model_impl + hf_runner, + vllm_runner, + example_prompts, + model, + num_fused=num_fused, + model_impl=model_impl, ) diff --git a/tests/multimodal/media/test_unprocessable_entity_error.py b/tests/multimodal/media/test_unprocessable_entity_error.py new file mode 100644 index 00000000000..8be70383b8a --- /dev/null +++ b/tests/multimodal/media/test_unprocessable_entity_error.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Tests for VLLMUnprocessableEntityError and media fetch error handling. + +Verifies that unprocessable image URLs (404, 403, DNS failures, etc.) return +HTTP 422 instead of 500. +""" + +from http import HTTPStatus +from unittest.mock import AsyncMock, MagicMock, patch + +import aiohttp +import pytest + +from vllm.entrypoints.serve.utils.error_response import create_error_response +from vllm.exceptions import VLLMUnprocessableEntityError +from vllm.multimodal.media import MediaConnector + + +class TestVLLMUnprocessableEntityError: + """Tests for VLLMUnprocessableEntityError exception.""" + + def test_creation(self): + exc = VLLMUnprocessableEntityError("Test error") + assert str(exc) == "Test error" + assert exc.parameter is None + + def test_creation_with_parameter_and_value(self): + exc = VLLMUnprocessableEntityError( + "Test error", + parameter="image_url", + value="https://example.com/image.jpg", + ) + assert "parameter=image_url" in str(exc) + assert "value=https://example.com/image.jpg" in str(exc) + + def test_is_value_error_subclass(self): + exc = VLLMUnprocessableEntityError("Test") + assert isinstance(exc, ValueError) + + +class TestMediaConnectorErrorHandling: + """Tests for MediaConnector error handling.""" + + @pytest.mark.asyncio + async def test_fetch_image_async_404(self): + connector = MediaConnector() + + with patch.object( + connector.connection, "async_get_bytes", new_callable=AsyncMock + ) as mock_get: + mock_get.side_effect = aiohttp.ClientResponseError( + request_info=MagicMock(), + history=(), + status=404, + message="Not Found", + ) + + with pytest.raises(VLLMUnprocessableEntityError) as exc_info: + await connector.fetch_image_async("https://example.com/missing.jpg") + + assert exc_info.value.parameter == "image_url" + + @pytest.mark.asyncio + async def test_fetch_image_async_dns_error(self): + """DNS errors are transient and should remain as-is for retry.""" + connector = MediaConnector() + + with patch.object( + connector.connection, "async_get_bytes", new_callable=AsyncMock + ) as mock_get: + mock_get.side_effect = aiohttp.ClientConnectorDNSError( + connection_key=MagicMock(), + os_error=MagicMock(), + ) + + with pytest.raises(aiohttp.ClientConnectorDNSError) as exc_info: + await connector.fetch_image_async( + "https://nonexistent.example/image.jpg" + ) + + assert isinstance(exc_info.value, aiohttp.ClientConnectorDNSError) + + @pytest.mark.asyncio + async def test_fetch_image_async_500_preserved(self): + """5xx errors should remain as server errors.""" + connector = MediaConnector() + + with patch.object( + connector.connection, "async_get_bytes", new_callable=AsyncMock + ) as mock_get: + mock_get.side_effect = aiohttp.ClientResponseError( + request_info=MagicMock(), + history=(), + status=500, + message="Internal Server Error", + ) + + with pytest.raises(aiohttp.ClientResponseError) as exc_info: + await connector.fetch_image_async("https://example.com/image.jpg") + + assert exc_info.value.status == 500 + + def test_fetch_image_404(self): + connector = MediaConnector() + + with patch.object( + connector.connection, "get_bytes", new_callable=MagicMock + ) as mock_get: + mock_get.side_effect = aiohttp.ClientResponseError( + request_info=MagicMock(), + history=(), + status=404, + message="Not Found", + ) + + with pytest.raises(VLLMUnprocessableEntityError) as exc_info: + connector.fetch_image("https://example.com/missing.jpg") + + assert exc_info.value.parameter == "image_url" + + def test_fetch_image_connection_error(self): + """Connection errors are transient and should remain as-is for retry.""" + connector = MediaConnector() + + with patch.object( + connector.connection, "get_bytes", new_callable=MagicMock + ) as mock_get: + mock_get.side_effect = aiohttp.ClientConnectionError("Connection refused") + + with pytest.raises(aiohttp.ClientConnectionError) as exc_info: + connector.fetch_image("https://example.com/image.jpg") + + assert isinstance(exc_info.value, aiohttp.ClientConnectionError) + + +class TestErrorResponse: + """Tests for error response creation.""" + + def test_unprocessable_entity_returns_422(self): + exc = VLLMUnprocessableEntityError( + "Failed to fetch media from URL: Cannot connect", + parameter="image_url", + value="https://example.com/image.jpg", + ) + + response = create_error_response(exc) + + assert response.error.code == HTTPStatus.UNPROCESSABLE_ENTITY.value + assert response.error.type == "UnprocessableEntityError" + assert response.error.param == "image_url" + + def test_unprocessable_entity_message(self): + exc = VLLMUnprocessableEntityError("Test error message") + response = create_error_response(exc) + + assert response.error.message == "Test error message" + assert response.error.code == 422 diff --git a/tests/multimodal/media/test_video.py b/tests/multimodal/media/test_video.py index e4b3afff084..671abd7b077 100644 --- a/tests/multimodal/media/test_video.py +++ b/tests/multimodal/media/test_video.py @@ -16,7 +16,11 @@ from vllm.assets.video import ( video_to_pil_images_list, ) from vllm.multimodal.media import ImageMediaIO, VideoMediaIO -from vllm.multimodal.video import VIDEO_LOADER_REGISTRY, VideoLoader +from vllm.multimodal.video import ( + PYNVVIDEOCODEC_VIDEO_BACKEND, + VIDEO_LOADER_REGISTRY, + VideoLoader, +) from ..utils import cosine_similarity, create_video_from_image, normalize_image @@ -357,3 +361,95 @@ def test_load_base64_jpeg_raises_on_zero_num_frames(): with pytest.raises(ValueError, match="num_frames must be greater than 0 or -1"): videoio.load_base64("video/jpeg", data) + + +# --------------------------------------------------------------------------- +# GPU video backend policy tests +# --------------------------------------------------------------------------- + + +class TestMergeKwargsGpuBackendPolicy: + """Verify that merge_kwargs blocks request-level GPU backend selection + when the static (engine-level) config did not configure that backend.""" + + def test_pynvvideocodec_requires_gpu(self): + assert VIDEO_LOADER_REGISTRY.backend_requires_gpu(PYNVVIDEOCODEC_VIDEO_BACKEND) + + def test_strips_video_backend_pynv_when_not_static(self): + result = VideoMediaIO.merge_kwargs( + default_kwargs=None, + runtime_kwargs={"video_backend": "pynvvideocodec"}, + ) + assert "video_backend" not in result + + def test_strips_backend_pynv_when_not_static(self): + result = VideoMediaIO.merge_kwargs( + default_kwargs={"num_frames": 16}, + runtime_kwargs={"backend": "pynvvideocodec"}, + ) + assert result.get("backend") != "pynvvideocodec" + + def test_preserves_video_backend_pynv_when_static(self): + result = VideoMediaIO.merge_kwargs( + default_kwargs={"video_backend": "pynvvideocodec"}, + runtime_kwargs={"video_backend": "pynvvideocodec", "num_frames": 8}, + ) + assert result["video_backend"] == "pynvvideocodec" + assert result["num_frames"] == 8 + + def test_preserves_backend_pynv_when_static(self): + result = VideoMediaIO.merge_kwargs( + default_kwargs={"backend": "pynvvideocodec"}, + runtime_kwargs={"backend": "pynvvideocodec"}, + ) + assert result["backend"] == "pynvvideocodec" + + @pytest.mark.parametrize("backend", ["opencv", "pyav", "torchcodec"]) + def test_software_video_backend_passes_through(self, backend: str): + result = VideoMediaIO.merge_kwargs( + default_kwargs=None, + runtime_kwargs={"video_backend": backend}, + ) + assert result["video_backend"] == backend + + @pytest.mark.parametrize("backend", ["opencv", "pyav"]) + def test_software_codec_backend_passes_through(self, backend: str): + result = VideoMediaIO.merge_kwargs( + default_kwargs=None, + runtime_kwargs={"backend": backend}, + ) + assert result["backend"] == backend + + def test_strips_both_keys_independently(self): + result = VideoMediaIO.merge_kwargs( + default_kwargs=None, + runtime_kwargs={ + "video_backend": "pynvvideocodec", + "backend": "pynvvideocodec", + "num_frames": 4, + }, + ) + assert "video_backend" not in result + assert result.get("backend") != "pynvvideocodec" + assert result["num_frames"] == 4 + + def test_other_kwargs_preserved_when_gpu_backend_stripped(self): + result = VideoMediaIO.merge_kwargs( + default_kwargs={"fps": 2}, + runtime_kwargs={ + "video_backend": "pynvvideocodec", + "num_frames": 16, + }, + ) + assert "video_backend" not in result + assert result["num_frames"] == 16 + + def test_static_pynv_with_different_runtime_gpu_backend(self): + """If static sets pynv via video_backend but runtime tries to set it + via the codec-level 'backend' key (without a static match), strip it.""" + result = VideoMediaIO.merge_kwargs( + default_kwargs={"video_backend": "pynvvideocodec"}, + runtime_kwargs={"backend": "pynvvideocodec"}, + ) + assert result.get("backend") != "pynvvideocodec" + assert result["video_backend"] == "pynvvideocodec" diff --git a/tests/multimodal/test_video.py b/tests/multimodal/test_video.py index 6fccc926a21..6aeb7dc486e 100644 --- a/tests/multimodal/test_video.py +++ b/tests/multimodal/test_video.py @@ -768,6 +768,118 @@ def test_pyav_backend_returns_target_frames_not_keyframes(): ) +# ============================================================================ +# TorchCodec Backend Tests +# ============================================================================ + + +def test_torchcodec_backend_loads_frames( + dummy_video_path, monkeypatch: pytest.MonkeyPatch +): + """Test that the torchcodec codec backend can load frames.""" + pytest.importorskip("torchcodec") + with monkeypatch.context() as m: + m.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv") + + with open(dummy_video_path, "rb") as f: + video_data = f.read() + + loader = VIDEO_LOADER_REGISTRY.load("opencv") + frames, metadata = loader.load_bytes( + video_data, num_frames=8, backend="torchcodec" + ) + + assert frames.ndim == 4 + assert frames.shape[3] == 3 # RGB + assert frames.shape[0] == 8 + assert frames.shape[0] == len(metadata["frames_indices"]) + assert metadata["video_backend"] == "torchcodec" + assert "total_num_frames" in metadata + assert "fps" in metadata + assert "duration" in metadata + + +def test_torchcodec_dynamic_backend_loads_frames( + dummy_video_path, monkeypatch: pytest.MonkeyPatch +): + """Test that the torchcodec codec with dynamic sampling can load frames.""" + pytest.importorskip("torchcodec") + with monkeypatch.context() as m: + m.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv_dynamic") + + with open(dummy_video_path, "rb") as f: + video_data = f.read() + + loader = VIDEO_LOADER_REGISTRY.load("opencv_dynamic") + frames, metadata = loader.load_bytes( + video_data, fps=2, max_duration=10, backend="torchcodec" + ) + + assert frames.ndim == 4 + assert frames.shape[3] == 3 # RGB + assert frames.shape[0] > 0 + assert frames.shape[0] == len(metadata["frames_indices"]) + assert metadata["video_backend"] == "torchcodec_dynamic" + + +def test_torchcodec_backend_rejects_frame_recovery(dummy_video_path): + """frame_recovery is OpenCV-only; torchcodec must reject it.""" + pytest.importorskip("torchcodec") + with open(dummy_video_path, "rb") as f: + video_data = f.read() + + loader = VIDEO_LOADER_REGISTRY.load("opencv") + with pytest.raises(AssertionError): + loader.load_bytes( + video_data, num_frames=8, backend="torchcodec", frame_recovery=True + ) + + +def test_torchcodec_backend_returns_target_frames_not_keyframes(): + """Regression test: torchcodec must return the requested frames, not the + GOP keyframe they seek back to. + + Mirrors ``test_pyav_backend_returns_target_frames_not_keyframes``: a long + GOP (single keyframe at frame 0) with a per-frame green-channel marker. + With ``seek_mode="exact"`` torchcodec resolves each index to the exact + frame, so the returned markers must be distinct, ordered, and match the + requested indices. + """ + pytest.importorskip("torchcodec") + num_frames = 50 + num_sampled = 4 + height, width = 64, 64 + + video_bytes = create_long_gop_video( + num_frames=num_frames, width=width, height=height + ) + + loader = VIDEO_LOADER_REGISTRY.load("opencv") + frames, metadata = loader.load_bytes( + video_bytes, num_frames=num_sampled, backend="torchcodec" + ) + assert frames.shape == (num_sampled, height, width, 3) + + requested = list(metadata["frames_indices"]) + assert len(requested) == num_sampled + + actual = [int(f[height // 2, width // 2, 1]) for f in frames] + + assert len(set(actual)) == num_sampled, ( + f"torchcodec returned only {len(set(actual))} distinct frames for " + f"{num_sampled} requested indices: markers={actual}, " + f"requested={requested}. Keyframe-snap regression." + ) + + assert actual == sorted(actual), f"Returned frames out of order: markers={actual}" + + for marker, want_idx in zip(actual, requested): + assert abs(marker - want_idx) <= 10, ( + f"Frame mismatch: requested index {want_idx}, " + f"got marker {marker} (tolerance ±10)" + ) + + @pytest.mark.parametrize( "loader_key, kwargs, expected_num_frames", [ @@ -854,6 +966,42 @@ def test_pyav_backend_returns_target_frames_not_keyframes(): 120, id="glm46v-pyav-60s", ), + # uniform sampling + torchcodec codec (same frame counts as opencv) + pytest.param( + "opencv", + {"num_frames": 32, "backend": "torchcodec"}, + 32, + id="torchcodec-num_frames", + ), + pytest.param( + "opencv", {"fps": 2, "backend": "torchcodec"}, 120, id="torchcodec-fps" + ), + pytest.param( + "opencv", + {"num_frames": 500, "fps": 2, "backend": "torchcodec"}, + 120, + id="torchcodec-num_frames_wins_fps", + ), + # dynamic sampling + torchcodec codec + pytest.param( + "opencv_dynamic", + {"fps": 1, "max_duration": 60, "backend": "torchcodec"}, + 60, + id="torchcodec_dynamic-within_max_duration", + ), + pytest.param( + "opencv_dynamic", + {"fps": 2, "max_duration": 30, "backend": "torchcodec"}, + 60, + id="torchcodec_dynamic-exceeds_max_duration", + ), + # glm46v dynamic FPS + torchcodec codec + pytest.param( + "glm46v", + {"backend": "torchcodec"}, + 120, + id="glm46v-torchcodec-60s", + ), ], ) def test_video_loader_frames_sampling( @@ -864,6 +1012,8 @@ def test_video_loader_frames_sampling( expected_num_frames: int, ): """Test video loader frames sampling functionality.""" + if kwargs.get("backend") == "torchcodec": + pytest.importorskip("torchcodec") monkeypatch.setenv("VLLM_VIDEO_LOADER_BACKEND", loader_key) loader = VIDEO_LOADER_REGISTRY.load(loader_key) diff --git a/tests/quantization/test_auto_round.py b/tests/quantization/test_auto_round.py index f5a38ddb51d..2cd70e4cf0f 100644 --- a/tests/quantization/test_auto_round.py +++ b/tests/quantization/test_auto_round.py @@ -61,7 +61,7 @@ MODELS = [ ) @pytest.mark.parametrize("model", MODELS) def test_auto_round_model(vllm_runner, model): - with vllm_runner(model, enforce_eager=True) as llm: + with vllm_runner(model) as llm: output = llm.generate_greedy(["The capital of France is"], max_tokens=8) assert output @@ -336,7 +336,7 @@ def test_wna16_xpu_prefers_ark_when_available(monkeypatch) -> None: monkeypatch.setattr(current_platform, "is_xpu", lambda: True) monkeypatch.setattr(current_platform, "is_cpu", lambda: False) monkeypatch.setattr( - "vllm.model_executor.layers.quantization.inc.schemes.inc_wna16_linear.get_ark_state", + "vllm.model_executor.layers.quantization.inc.schemes.inc_ark_ops.get_ark_state", lambda: (True, None, object(), DummyQuantLinear), ) @@ -355,7 +355,7 @@ def test_wna16_xpu_falls_back_when_ark_unavailable(monkeypatch) -> None: monkeypatch.setattr(current_platform, "is_xpu", lambda: True) monkeypatch.setattr(current_platform, "is_cpu", lambda: False) monkeypatch.setattr( - "vllm.model_executor.layers.quantization.inc.schemes.inc_wna16_linear.get_ark_state", + "vllm.model_executor.layers.quantization.inc.schemes.inc_ark_ops.get_ark_state", lambda: (False, "missing", None, None), ) @@ -377,7 +377,7 @@ def test_wna16_cpu_gptq_prefers_ark_when_available(monkeypatch) -> None: monkeypatch.setattr(current_platform, "is_xpu", lambda: False) monkeypatch.setattr(current_platform, "is_cpu", lambda: True) monkeypatch.setattr( - "vllm.model_executor.layers.quantization.inc.schemes.inc_wna16_linear.get_ark_state", + "vllm.model_executor.layers.quantization.inc.schemes.inc_ark_ops.get_ark_state", lambda: (True, None, object(), DummyQuantLinear), ) @@ -398,7 +398,7 @@ def test_wna16_cpu_gptq_raises_when_ark_and_marlin_unavailable( monkeypatch.setattr(current_platform, "is_xpu", lambda: False) monkeypatch.setattr(current_platform, "is_cpu", lambda: True) monkeypatch.setattr( - "vllm.model_executor.layers.quantization.inc.schemes.inc_wna16_linear.get_ark_state", + "vllm.model_executor.layers.quantization.inc.schemes.inc_ark_ops.get_ark_state", lambda: (False, "missing", None, None), ) monkeypatch.setattr( diff --git a/tests/test_config.py b/tests/test_config.py index 3837057658b..1e93b610da5 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -649,6 +649,30 @@ def test_nested_hf_overrides(): assert model_config.hf_config.vision_config.hidden_size == 512 +def test_model_class_overrides_registers_target(): + """`model_class_overrides` redirects an architecture to a custom class.""" + from vllm.model_executor.models import ModelRegistry + + arch = "_TestModelClassOverrideArch" + target = "vllm.model_executor.models.llama:LlamaForCausalLM" + assert arch not in ModelRegistry.models + + model_config = ModelConfig( + "facebook/opt-125m", + model_class_overrides={arch: target}, + ) + try: + # Accessing `.registry` is the chokepoint that applies the overrides; + # it has already run during construction. + registered = model_config.registry.models[arch] + assert registered.module_name == "vllm.model_executor.models.llama" + assert registered.class_name == "LlamaForCausalLM" + # Idempotent: a second access does not re-register or error out. + assert model_config.registry.models[arch] is registered + finally: + ModelRegistry.models.pop(arch, None) + + @pytest.mark.skipif( current_platform.is_rocm(), reason="Encoder Decoder models not supported on ROCm." ) diff --git a/tests/tool_parsers/test_glm47_moe_tool_parser.py b/tests/tool_parsers/test_glm47_moe_tool_parser.py index c9767f6f62f..224fea08d74 100644 --- a/tests/tool_parsers/test_glm47_moe_tool_parser.py +++ b/tests/tool_parsers/test_glm47_moe_tool_parser.py @@ -7,12 +7,16 @@ import json from unittest.mock import Mock import pytest +from openai.types.responses import ResponseFunctionToolCall from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ChatCompletionToolsParam, FunctionDefinition, ) +from vllm.entrypoints.openai.engine.protocol import FunctionCall +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.entrypoints.openai.responses.utils import build_response_output_items from vllm.tokenizers import get_tokenizer from vllm.tool_parsers.glm47_moe_tool_parser import Glm47MoeModelToolParser @@ -58,7 +62,69 @@ def mock_request(sample_tools) -> ChatCompletionRequest: return request +@pytest.fixture +def namespace_tool_request() -> ResponsesRequest: + return ResponsesRequest.model_validate( + { + "input": "hi", + "tools": [ + { + "type": "namespace", + "name": "mcp__computer_use", + "description": "Computer use tools.", + "tools": [ + { + "type": "function", + "name": "get_app_state", + "description": "Get app state.", + "parameters": { + "type": "object", + "properties": { + "app": {"type": "string"}, + }, + }, + } + ], + } + ], + } + ) + + class TestGlm47ExtractToolCalls: + def test_namespace_tool_call_round_trip_to_responses_output( + self, glm47_tokenizer, namespace_tool_request + ): + parser = Glm47MoeModelToolParser( + glm47_tokenizer, tools=namespace_tool_request.tools + ) + out = ( + "mcp__computer_use__get_app_state" + "app" + "Google Chrome" + "" + ) + + result = parser.extract_tool_calls(out, request=namespace_tool_request) + + assert result.tools_called + tool_call = result.tool_calls[0].function + assert tool_call == FunctionCall( + name="mcp__computer_use__get_app_state", + arguments='{"app": "Google Chrome"}', + ) + + output_items = build_response_output_items( + reasoning=None, + content=None, + tool_calls=[tool_call], + tools=namespace_tool_request.tools, + ) + output_tool_call = output_items[0] + assert isinstance(output_tool_call, ResponseFunctionToolCall) + assert output_tool_call.name == "get_app_state" + assert output_tool_call.namespace == "mcp__computer_use" + def test_no_tool_call(self, glm47_tool_parser, mock_request): out = "This is a plain response." r = glm47_tool_parser.extract_tool_calls(out, request=mock_request) diff --git a/tests/tool_parsers/test_granite_tool_parser.py b/tests/tool_parsers/test_granite_tool_parser.py index 2046c11c5d2..af3386112f5 100644 --- a/tests/tool_parsers/test_granite_tool_parser.py +++ b/tests/tool_parsers/test_granite_tool_parser.py @@ -2,13 +2,21 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json + import pytest from tests.tool_parsers.common_tests import ( ToolParserTestConfig, ToolParserTests, ) -from tests.tool_parsers.utils import run_tool_extraction +from tests.tool_parsers.utils import ( + run_tool_extraction, + run_tool_extraction_streaming, + split_string_into_token_deltas, +) +from vllm.tokenizers import get_tokenizer +from vllm.tool_parsers.granite_tool_parser import GraniteToolParser class TestGraniteToolParser(ToolParserTests): @@ -116,3 +124,38 @@ I'll get that information.""", f"Expected 1 tool call from string format, got {len(tool_calls)}" ) assert tool_calls[0].function.name == "get_weather" + + +# granite emits arguments before name and its own tokenizer (not gpt2) is used +# here so the token boundaries match production; get_tokenizer only fetches the +# small tokenizer files, not the model weights. +@pytest.fixture(scope="module") +def granite_tokenizer(): + return get_tokenizer(tokenizer_name="ibm-granite/granite-3.1-8b-instruct") + + +@pytest.mark.parametrize("chunk_size", [2, 3, 4, 5]) +def test_streaming_parallel_calls_batched_deltas(granite_tokenizer, chunk_size): + """A batched delta (multiple tokens) spanning the boundary between two + parallel calls must not drop the first call's name. granite streams + arguments before name, so the name only completes as the next call appears. + """ + parser = GraniteToolParser(granite_tokenizer) + model_output = ( + '<|tool_call|> [{"arguments": {"city": "Tokyo"}, "name": "get_weather"}, ' + '{"arguments": {"timezone": "Asia/Tokyo"}, "name": "get_time"}]' + ) + token_deltas = split_string_into_token_deltas(granite_tokenizer, model_output) + batched = [ + "".join(token_deltas[i : i + chunk_size]) + for i in range(0, len(token_deltas), chunk_size) + ] + reconstructor = run_tool_extraction_streaming( + parser, batched, assert_one_tool_per_delta=False + ) + names = [tc.function.name for tc in reconstructor.tool_calls] + assert names == ["get_weather", "get_time"] + # trailing args of the final call are flushed by the serving layer + assert json.loads(reconstructor.tool_calls[0].function.arguments) == { + "city": "Tokyo" + } diff --git a/tests/tool_parsers/test_mistral_tool_parser.py b/tests/tool_parsers/test_mistral_tool_parser.py index 03a10ef0991..4f75e2f576f 100644 --- a/tests/tool_parsers/test_mistral_tool_parser.py +++ b/tests/tool_parsers/test_mistral_tool_parser.py @@ -144,6 +144,7 @@ def stream_delta_message_generator( mistral_tokenizer: TokenizerLike, model_output: str | None, tools: list[tuple[str, str]] | None, + chunk_size: int = 1, ) -> Generator[DeltaMessage, None, None]: if ( isinstance(mistral_tokenizer, MistralTokenizer) @@ -182,15 +183,13 @@ def stream_delta_message_generator( previous_tokens = None prefix_offset = 0 read_offset = 0 + pending_text = "" + pending_token_ids: list[int] = [] for i, delta_token in enumerate(all_token_ids): - delta_token_ids = [delta_token] - previous_token_ids = all_token_ids[:i] - current_token_ids = all_token_ids[: i + 1] - (new_tokens, delta_text, new_prefix_offset, new_read_offset) = ( detokenize_incrementally( tokenizer=mistral_tokenizer, - all_input_ids=current_token_ids, + all_input_ids=all_token_ids[: i + 1], prev_tokens=previous_tokens, prefix_offset=prefix_offset, read_offset=read_offset, @@ -198,27 +197,39 @@ def stream_delta_message_generator( spaces_between_special_tokens=True, ) ) + previous_tokens = ( + previous_tokens + new_tokens if previous_tokens else new_tokens + ) + prefix_offset = new_prefix_offset + read_offset = new_read_offset - current_text = previous_text + delta_text + # Buffer tokens so each streamed delta can carry ``chunk_size`` tokens, + # reproducing the multi-token deltas produced by async scheduling / + # stream_interval > 1. + pending_text += delta_text + pending_token_ids.append(delta_token) + if len(pending_token_ids) < chunk_size and i != len(all_token_ids) - 1: + continue + + previous_token_ids = all_token_ids[: i + 1 - len(pending_token_ids)] + current_token_ids = all_token_ids[: i + 1] + current_text = previous_text + pending_text delta_message = mistral_tool_parser.extract_tool_calls_streaming( previous_text, current_text, - delta_text, + pending_text, previous_token_ids, current_token_ids, - delta_token_ids, + pending_token_ids, request=_DUMMY_REQUEST, ) if delta_message: yield delta_message previous_text = current_text - previous_tokens = ( - previous_tokens + new_tokens if previous_tokens else new_tokens - ) - prefix_offset = new_prefix_offset - read_offset = new_read_offset + pending_text = "" + pending_token_ids = [] @pytest.mark.parametrize( @@ -1572,3 +1583,39 @@ def test_grammar_from_tool_parser_set_by_adjust_request( request = _make_request() result = mistral_tool_parser.adjust_request(request) assert result._grammar_from_tool_parser is True + + +@pytest.mark.parametrize("chunk_size", [2, 3, 4, 5]) +def test_streaming_pre_v11_parallel_calls_batched_deltas( + mistral_pre_v11_tool_parser, mistral_pre_v11_tokenizer, chunk_size +): + """A batched delta spanning the boundary between two parallel calls must + keep them on distinct indices (the bug collapsed both onto index 0).""" + model_output = ( + '[TOOL_CALLS] [{"name": "add", "arguments": {"a": 3.5, "b": 4}}, ' + '{"name": "get_current_weather", "arguments": ' + '{"city": "San Francisco", "state": "CA", "unit": "celsius"}}]' + ) + names: list[str] = [] + args: list[str] = [] + idx = -1 + for delta_message in stream_delta_message_generator( + mistral_pre_v11_tool_parser, + mistral_pre_v11_tokenizer, + model_output, + tools=None, + chunk_size=chunk_size, + ): + for tool_call in delta_message.tool_calls or []: + if tool_call.index != idx: + idx = tool_call.index + args.append("") + if tool_call.function and tool_call.function.name: + names.append(tool_call.function.name) + if tool_call.function and tool_call.function.arguments: + args[tool_call.index] += tool_call.function.arguments + + assert names == ["add", "get_current_weather"] + assert len(args) == 2 + # trailing args of the final call are flushed by the serving layer + assert json.loads(args[0]) == {"a": 3.5, "b": 4} diff --git a/tests/tool_use/test_parallel_tool_calls.py b/tests/tool_use/test_parallel_tool_calls.py index 0f7f6893162..4cfd165f1a8 100644 --- a/tests/tool_use/test_parallel_tool_calls.py +++ b/tests/tool_use/test_parallel_tool_calls.py @@ -115,14 +115,12 @@ async def test_parallel_tool_calls( assert not role_name or role_name == "assistant" role_name = "assistant" - # if a tool call is streamed make sure there's exactly one - # (based on the request parameters + # a chunk may carry >1 tool-call delta at a parallel-call boundary streamed_tool_calls = chunk.choices[0].delta.tool_calls - if streamed_tool_calls and len(streamed_tool_calls) > 0: - # make sure only one diff is present - correct even for parallel - assert len(streamed_tool_calls) == 1 - tool_call = streamed_tool_calls[0] + for tool_call in streamed_tool_calls or []: + # deltas arrive in non-decreasing index order + assert tool_call.index >= tool_call_idx # if a new tool is being called, set up empty arguments if tool_call.index != tool_call_idx: diff --git a/tests/utils.py b/tests/utils.py index 08579f99e4d..2a3bdb91fe0 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1459,6 +1459,46 @@ def multi_process_parallel( ray.shutdown() +def assert_rocm_custom_allreduce_backend_state( + use_aiter_custom_ar: bool, + quick_reduce_quantization: str, +) -> None: + from vllm.distributed.parallel_state import get_tp_group + + device_communicator = get_tp_group().device_communicator + aiter_ar_comm = device_communicator.aiter_ar_comm + if use_aiter_custom_ar: + assert aiter_ar_comm is not None, "AITER CustomAllreduce was not initialized." + assert not aiter_ar_comm.disabled, "AITER CustomAllreduce is disabled." + assert device_communicator.ca_comm is None, ( + "vLLM CustomAllreduce should not be initialized when AITER CA is used." + ) + else: + assert aiter_ar_comm is None, ( + "AITER CustomAllreduce should not be initialized when disabled." + ) + assert device_communicator.ca_comm is not None, ( + "vLLM CustomAllreduce should be initialized when AITER CA is disabled." + ) + + qr_comm = device_communicator.qr_comm + assert qr_comm is not None, "QuickReduce communicator was not initialized." + if quick_reduce_quantization == "NONE": + assert qr_comm.disabled, "QuickReduce should be disabled." + else: + assert not qr_comm.disabled, "QuickReduce should be enabled." + + +def assert_rocm_custom_allreduce_backend_state_on_worker( + _worker, + use_aiter_custom_ar: bool, + quick_reduce_quantization: str, +) -> None: + assert_rocm_custom_allreduce_backend_state( + use_aiter_custom_ar, quick_reduce_quantization + ) + + @contextmanager def error_on_warning(category: type[Warning] = Warning): """ diff --git a/tests/v1/attention/test_linear_attention_metadata_builder.py b/tests/v1/attention/test_linear_attention_metadata_builder.py new file mode 100644 index 00000000000..3ef811b3a66 --- /dev/null +++ b/tests/v1/attention/test_linear_attention_metadata_builder.py @@ -0,0 +1,188 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from tests.v1.attention.utils import ( + BatchSpec, + create_common_attn_metadata, + create_vllm_config, +) +from vllm.config import CUDAGraphMode, SpeculativeConfig +from vllm.v1.attention.backend import AttentionCGSupport +from vllm.v1.attention.backends.linear_attn import ( + BailingLinearAttentionMetadataBuilder, + LinearAttentionMetadataBuilder, +) +from vllm.v1.attention.backends.utils import PAD_SLOT_ID +from vllm.v1.kv_cache_interface import MambaSpec + +BLOCK_SIZE = 16 +DEVICE = torch.device("cpu") + + +def _create_mamba_spec(num_speculative_blocks: int = 1) -> MambaSpec: + return MambaSpec( + block_size=BLOCK_SIZE, + shapes=((16, 64),), + dtypes=(torch.float16,), + num_speculative_blocks=num_speculative_blocks, + ) + + +def test_bailing_linear_attention_reports_uniform_batch_cudagraph_support(): + vllm_config = create_vllm_config( + hf_config_override={ + "architectures": ["BailingMoeV2_5ForCausalLM"], + "model_type": "bailing_hybrid", + } + ) + + support = BailingLinearAttentionMetadataBuilder.get_cudagraph_support( + vllm_config, _create_mamba_spec() + ) + + assert support == AttentionCGSupport.UNIFORM_BATCH + + +def test_non_bailing_linear_attention_keeps_single_token_cudagraph_support(): + vllm_config = create_vllm_config( + hf_config_override={ + "architectures": ["MiniMaxText01ForCausalLM"], + "model_type": "minimax_text_01", + } + ) + + support = LinearAttentionMetadataBuilder.get_cudagraph_support( + vllm_config, _create_mamba_spec() + ) + + assert support == AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE + + +def test_linear_attention_spec_decode_full_graph_metadata_pads_cache_slots(): + vllm_config = create_vllm_config( + hf_config_override={ + "architectures": ["BailingMoeV2_5ForCausalLM"], + "model_type": "bailing_hybrid", + } + ) + vllm_config.speculative_config = SpeculativeConfig( + method="ngram", + num_speculative_tokens=1, + ) + vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY + + builder = BailingLinearAttentionMetadataBuilder( + kv_cache_spec=_create_mamba_spec(), + layer_names=["model.layers.0.self_attn"], + vllm_config=vllm_config, + device=DEVICE, + ) + + common = create_common_attn_metadata( + BatchSpec(seq_lens=[20, 20, 0], query_lens=[2, 2, 0]), + BLOCK_SIZE, + DEVICE, + ) + common.block_table_tensor[2].fill_(-1) + + metadata = builder.build( + common_prefix_len=0, + common_attn_metadata=common, + num_accepted_tokens=torch.tensor([1, 2, 1], dtype=torch.int32), + ) + + assert metadata.num_decodes == 3 + assert metadata.num_prefills == 0 + assert metadata.num_decode_tokens == 4 + assert metadata.state_indices_tensor_d is not None + assert metadata.state_indices_tensor_d.shape == (3, 2) + assert torch.equal( + metadata.state_indices_tensor_d[2], + torch.full((2,), PAD_SLOT_ID, dtype=torch.int32), + ) + assert torch.equal( + metadata.state_indices_tensor[2], + torch.tensor(PAD_SLOT_ID, dtype=torch.int32), + ) + assert metadata.query_start_loc_d is not None + assert metadata.query_start_loc_d.tolist() == [0, 2, 4, 4] + assert metadata.num_accepted_tokens is not None + assert metadata.num_accepted_tokens.tolist() == [1, 2, 1] + + +def test_linear_attention_full_graph_metadata_uses_stable_decode_buffers(): + vllm_config = create_vllm_config( + hf_config_override={ + "architectures": ["BailingMoeV2_5ForCausalLM"], + "model_type": "bailing_hybrid", + } + ) + vllm_config.speculative_config = SpeculativeConfig( + method="ngram", + num_speculative_tokens=1, + ) + vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY + + builder = BailingLinearAttentionMetadataBuilder( + kv_cache_spec=_create_mamba_spec(), + layer_names=["model.layers.0.self_attn"], + vllm_config=vllm_config, + device=DEVICE, + ) + + common = create_common_attn_metadata( + BatchSpec(seq_lens=[20, 20, 0], query_lens=[2, 2, 0]), + BLOCK_SIZE, + DEVICE, + arange_block_indices=True, + ) + common.block_table_tensor = torch.tensor( + [[10, 11], [12, 13], [-1, -1]], + dtype=torch.int32, + device=DEVICE, + ) + + first = builder.build( + common_prefix_len=0, + common_attn_metadata=common, + num_accepted_tokens=torch.tensor([1, 2, 1], dtype=torch.int32), + ) + assert first.state_indices_tensor_d is not None + assert first.query_start_loc_d is not None + assert first.num_accepted_tokens is not None + state_ptr = first.state_indices_tensor_d.data_ptr() + query_ptr = first.query_start_loc_d.data_ptr() + accepted_ptr = first.num_accepted_tokens.data_ptr() + + common2 = create_common_attn_metadata( + BatchSpec(seq_lens=[36, 0, 0], query_lens=[2, 0, 0]), + BLOCK_SIZE, + DEVICE, + arange_block_indices=True, + ) + common2.block_table_tensor = torch.tensor( + [[20, 21], [-1, -1], [-1, -1]], + dtype=torch.int32, + device=DEVICE, + ) + second = builder.build( + common_prefix_len=0, + common_attn_metadata=common2, + num_accepted_tokens=torch.tensor([2, 1, 1], dtype=torch.int32), + ) + + assert second.state_indices_tensor_d is not None + assert second.query_start_loc_d is not None + assert second.num_accepted_tokens is not None + assert second.state_indices_tensor_d.data_ptr() == state_ptr + assert second.query_start_loc_d.data_ptr() == query_ptr + assert second.num_accepted_tokens.data_ptr() == accepted_ptr + assert second.state_indices_tensor_d.tolist() == [ + [20, 21], + [PAD_SLOT_ID, PAD_SLOT_ID], + [PAD_SLOT_ID, PAD_SLOT_ID], + ] + assert second.query_start_loc_d.tolist() == [0, 2, 2, 2] + assert second.num_accepted_tokens.tolist() == [2, 1, 1] diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index 8253a8422e2..9b6f6458961 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -324,3 +324,48 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): assert request.status == RequestStatus.FINISHED_ERROR assert request.request_id not in scheduler.requests assert not scheduler.running + + +def test_no_placeholder_underflow_on_discarded_spec_frame(): + num_spec = 5 + scheduler = create_scheduler( + async_scheduling=True, + num_speculative_tokens=num_spec, + speculative_method="ngram_gpu", + ) + req = create_requests(num_requests=1, max_tokens=20)[0] + req.num_computed_tokens = req.num_tokens + scheduler.requests[req.request_id] = req + scheduler.running.append(req) + req.status = RequestStatus.RUNNING + + req.num_output_placeholders = 1 + req.async_tokens_to_discard = num_spec + computed_before = req.num_computed_tokens + + scheduler_output = SchedulerOutput( + scheduled_new_reqs=[], + scheduled_cached_reqs=CachedRequestData.make_empty(), + num_scheduled_tokens={req.request_id: num_spec + 1}, + total_num_scheduled_tokens=num_spec + 1, + scheduled_encoder_inputs={}, + scheduled_spec_decode_tokens={req.request_id: [10] * num_spec}, + num_common_prefix_blocks=[], + finished_req_ids=set(), + free_encoder_mm_hashes=[], + ) + model_runner_output = ModelRunnerOutput( + req_ids=[req.request_id], + req_id_to_index={req.request_id: 0}, + sampled_token_ids=[[999]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + + scheduler.update_from_output(scheduler_output, model_runner_output) + + assert req.num_output_placeholders == 1 + assert req.num_computed_tokens == computed_before + assert req.async_tokens_to_discard == num_spec - 1 + assert req.status == RequestStatus.RUNNING diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index 2450b23669a..3a375ec7574 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -54,6 +54,7 @@ def create_scheduler( block_size: int = 16, max_model_len: int | None = None, num_speculative_tokens: int | None = None, + speculative_method: str | None = None, skip_tokenizer_init: bool = False, async_scheduling: bool = False, pipeline_parallel_size: int = 1, @@ -126,9 +127,14 @@ def create_scheduler( speculative_config: SpeculativeConfig | None = None if num_speculative_tokens is not None: - speculative_config = SpeculativeConfig( + spec_kwargs: dict = dict( model="ngram", num_speculative_tokens=num_speculative_tokens ) + if speculative_method is not None: + spec_kwargs["method"] = speculative_method + spec_kwargs["prompt_lookup_max"] = num_speculative_tokens + spec_kwargs["prompt_lookup_min"] = 1 + speculative_config = SpeculativeConfig(**spec_kwargs) ec_transfer_config = ( ECTransferConfig( diff --git a/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py b/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py index 4bb8d63a8a2..11f77492d2e 100644 --- a/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py +++ b/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py @@ -45,7 +45,9 @@ def test_prompts(): use_fork_for_test = ( - fork_new_process_for_each_test if not current_platform.is_rocm() else lambda x: x + fork_new_process_for_each_test + if not (current_platform.is_rocm() or current_platform.is_xpu()) + else lambda x: x ) diff --git a/tests/v1/e2e/general/test_rocm_aiter_custom_ar.py b/tests/v1/e2e/general/test_rocm_aiter_custom_ar.py new file mode 100644 index 00000000000..18e51584fa8 --- /dev/null +++ b/tests/v1/e2e/general/test_rocm_aiter_custom_ar.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm._aiter_ops import is_aiter_found, rocm_aiter_ops +from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode +from vllm.envs import disable_envs_cache +from vllm.platforms import current_platform + +from ....conftest import VllmRunner +from ....utils import ( + assert_rocm_custom_allreduce_backend_state_on_worker, + multi_gpu_test, +) + +PROMPTS = ["Hello, my name is", "The capital of France is"] + + +def _run_generation( + vllm_runner: type[VllmRunner], + monkeypatch: pytest.MonkeyPatch, + compilation_config: CompilationConfig, + *, + model: str, + max_tokens: int, + use_aiter_custom_ar: bool, + quick_reduce_quantization: str, +) -> list[tuple[list[int], str]]: + with monkeypatch.context() as m: + m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") + m.setenv("VLLM_ROCM_USE_AITER", "1") + m.setenv( + "VLLM_ROCM_USE_AITER_CUSTOM_AR", + "1" if use_aiter_custom_ar else "0", + ) + m.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", quick_reduce_quantization) + disable_envs_cache() + rocm_aiter_ops.refresh_env_variables() + + with vllm_runner( + model, + dtype="half", + tensor_parallel_size=2, + compilation_config=compilation_config, + max_model_len=256, + max_num_seqs=len(PROMPTS), + gpu_memory_utilization=0.7, + ) as llm: + llm.get_llm().collective_rpc( + assert_rocm_custom_allreduce_backend_state_on_worker, + args=(use_aiter_custom_ar, quick_reduce_quantization), + ) + + return llm.generate_greedy(PROMPTS, max_tokens) + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-only") +@pytest.mark.skipif(not is_aiter_found(), reason="AITER is not installed") +@multi_gpu_test(num_gpus=2) +@pytest.mark.parametrize( + "quick_reduce_quantization", + [ + pytest.param("FP", id="quick-reduce-on"), + pytest.param("NONE", id="quick-reduce-off"), + ], +) +@pytest.mark.parametrize( + "cudagraph_mode", + [ + pytest.param(CUDAGraphMode.NONE, id="cudagraph-none"), + pytest.param(CUDAGraphMode.FULL, id="cudagraph-full"), + ], +) +@pytest.mark.parametrize( + "model,max_tokens", + [ + pytest.param("facebook/opt-125m", 8, id="opt-125m"), + ], +) +def test_rocm_aiter_custom_ar_e2e( + vllm_runner: type[VllmRunner], + monkeypatch: pytest.MonkeyPatch, + cudagraph_mode: CUDAGraphMode, + quick_reduce_quantization: str, + model: str, + max_tokens: int, +): + compilation_mode = ( + CompilationMode.NONE + if cudagraph_mode == CUDAGraphMode.NONE + else CompilationMode.VLLM_COMPILE + ) + compilation_config = CompilationConfig( + mode=compilation_mode, + cudagraph_mode=cudagraph_mode, + ) + + baseline_generations = _run_generation( + vllm_runner, + monkeypatch, + compilation_config, + model=model, + max_tokens=max_tokens, + use_aiter_custom_ar=False, + quick_reduce_quantization=quick_reduce_quantization, + ) + aiter_custom_ar_generations = _run_generation( + vllm_runner, + monkeypatch, + compilation_config, + model=model, + max_tokens=max_tokens, + use_aiter_custom_ar=True, + quick_reduce_quantization=quick_reduce_quantization, + ) + + assert aiter_custom_ar_generations == baseline_generations diff --git a/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh b/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh index a34f07edc97..d3fee8d6ba5 100755 --- a/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh @@ -9,8 +9,10 @@ PREFILL_GPU_ID=${PREFILL_GPU_ID:-0} DECODE_GPU_ID=${DECODE_GPU_ID:-1} MODEL=${MODEL:-"ibm-granite/granite-4.0-h-tiny"} GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.8} +VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-} +ATTENTION_BACKEND=${ATTENTION_BACKEND:-FLASHINFER} -echo "Running Mamba prefix cache test (GPUs: P=$PREFILL_GPU_ID, D=$DECODE_GPU_ID, model=$MODEL)" +echo "Running Mamba prefix cache test (GPUs: P=$PREFILL_GPU_ID, D=$DECODE_GPU_ID, model=$MODEL, backend=$ATTENTION_BACKEND)" KV_CONFIG='{"kv_connector":"NixlConnector","kv_role":"kv_both"}' @@ -36,6 +38,14 @@ cleanup_instances() { cleanup_instances +EXTRA_ARGS=() +if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then + IFS=',' read -r -a EXTRA_ARGS <<< "$VLLM_SERVE_EXTRA_ARGS" +fi +if [[ -n "$ATTENTION_BACKEND" ]]; then + EXTRA_ARGS+=(--attention-backend "$ATTENTION_BACKEND") +fi + # Start prefill instance PREFILL_PORT=8001 CUDA_VISIBLE_DEVICES=$PREFILL_GPU_ID \ @@ -51,8 +61,8 @@ vllm serve $MODEL \ --trust-remote-code \ --enable-prefix-caching \ --mamba-cache-mode all \ - --attention-backend FLASHINFER \ - --kv-transfer-config "$KV_CONFIG" & + --kv-transfer-config "$KV_CONFIG" \ + "${EXTRA_ARGS[@]}" & # Start decode instance DECODE_PORT=8002 @@ -69,8 +79,8 @@ vllm serve $MODEL \ --trust-remote-code \ --enable-prefix-caching \ --mamba-cache-mode all \ - --attention-backend FLASHINFER \ - --kv-transfer-config "$KV_CONFIG" & + --kv-transfer-config "$KV_CONFIG" \ + "${EXTRA_ARGS[@]}" & echo "Waiting for prefill instance on port $PREFILL_PORT..." wait_for_server "$PREFILL_PORT" diff --git a/tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh index 2e71858983e..dae632dfce8 100755 --- a/tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh @@ -18,6 +18,7 @@ # Environment variables: # MODEL_NAMES - model to test (default: Qwen/Qwen3-0.6B) # GPU_MEMORY_UTILIZATION - GPU memory fraction (default: 0.6) +# ATTENTION_BACKEND - optional attention backend for vllm serve # VLLM_SERVE_EXTRA_ARGS - comma-separated extra args for vllm serve # SKIP_CROSS_LAYERS - set to 1 to skip the cross-layer layout test # SKIP_NORMAL_LAYOUT - set to 1 to skip the normal layout test @@ -34,9 +35,11 @@ fi GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.6} BLOCK_SIZE=${BLOCK_SIZE:-128} +ATTENTION_BACKEND=${ATTENTION_BACKEND:-} VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-} -GIT_ROOT=$(git rev-parse --show-toplevel) +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" +GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}" SMI_BIN=$(which nvidia-smi || which rocm-smi || echo "") # ── KV transfer configs ───────────────────────────────────────────────── @@ -139,6 +142,9 @@ run_tests_for_model() { BASE_CMD="${BASE_CMD} $arg" done fi + if [[ -n "$ATTENTION_BACKEND" ]]; then + BASE_CMD="${BASE_CMD} --attention-backend $ATTENTION_BACKEND" + fi eval "$BASE_CMD &" # ── Start decode instance ── @@ -161,6 +167,9 @@ run_tests_for_model() { BASE_CMD="${BASE_CMD} $arg" done fi + if [[ -n "$ATTENTION_BACKEND" ]]; then + BASE_CMD="${BASE_CMD} --attention-backend $ATTENTION_BACKEND" + fi eval "$BASE_CMD &" # ── Wait for servers ── diff --git a/tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh b/tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh index a80950b3413..de6c9abcc6b 100755 --- a/tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh @@ -19,6 +19,7 @@ # MODEL_NAMES - model to test (default: Qwen/Qwen3-0.6B) # KV_CACHE_MEMORY_BYTES - GPU KV cache size in bytes (default: 268435456 = 256 MiB) # BLOCK_SIZE - KV cache block size (default: 128) +# ATTENTION_BACKEND - optional attention backend for vllm serve # VLLM_SERVE_EXTRA_ARGS - comma-separated extra args for vllm serve set -xe @@ -34,9 +35,11 @@ fi KV_CACHE_MEMORY_BYTES=${KV_CACHE_MEMORY_BYTES:-268435456} # 256 MiB MAX_MODEL_LEN=${MAX_MODEL_LEN:-2048} BLOCK_SIZE=${BLOCK_SIZE:-128} +ATTENTION_BACKEND=${ATTENTION_BACKEND:-} VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-} -GIT_ROOT=$(git rev-parse --show-toplevel) +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" +GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}" # ── KV transfer config ────────────────────────────────────────────────── @@ -110,6 +113,9 @@ run_tests_for_model() { BASE_CMD="${BASE_CMD} $arg" done fi + if [[ -n "$ATTENTION_BACKEND" ]]; then + BASE_CMD="${BASE_CMD} --attention-backend $ATTENTION_BACKEND" + fi eval "$BASE_CMD &" # ── Start decode instance ── @@ -133,6 +139,9 @@ run_tests_for_model() { BASE_CMD="${BASE_CMD} $arg" done fi + if [[ -n "$ATTENTION_BACKEND" ]]; then + BASE_CMD="${BASE_CMD} --attention-backend $ATTENTION_BACKEND" + fi eval "$BASE_CMD &" # ── Wait for servers ── diff --git a/tests/v1/kv_connector/unit/test_hf3fs_connector.py b/tests/v1/kv_connector/unit/test_hf3fs_connector.py index cd525e23b14..94bb94c6fbd 100644 --- a/tests/v1/kv_connector/unit/test_hf3fs_connector.py +++ b/tests/v1/kv_connector/unit/test_hf3fs_connector.py @@ -33,7 +33,7 @@ def hf3fs_stats(): def _make_cuda_event(): """Return a real CUDA event when available, otherwise a MagicMock.""" if torch.cuda.is_available(): - return torch.Event() + return torch.cuda.Event() return MagicMock() diff --git a/tests/v1/worker/test_xpu_model_runner.py b/tests/v1/worker/test_xpu_model_runner.py new file mode 100644 index 00000000000..5ddf490c9b4 --- /dev/null +++ b/tests/v1/worker/test_xpu_model_runner.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for ``vllm.v1.worker.xpu_model_runner`` (XPU worker / CUDA shims).""" + +import pytest +import torch +from torch._dynamo.variables.torch import TorchInGraphFunctionVariable + +from vllm.v1.worker.xpu_model_runner import _torch_cuda_wrapper + +# XPU-only: needs distinct torch.cuda vs torch.xpu current_stream symbols. +pytestmark = pytest.mark.skipif( + not hasattr(torch, "xpu") or not hasattr(torch.xpu, "current_stream"), + reason="torch.xpu.current_stream is required", +) + + +# Child process: patched torch.cuda must not leak to other tests in the session. +@pytest.mark.forked +def test_torch_cuda_wrapper_allows_dynamo_handler_registration() -> None: + """Guard against XPU CUDA shim breaking Torch Dynamo during AOT compile. + + Before the fix, ``_torch_cuda_wrapper`` assigned + ``torch.cuda.current_stream = torch.xpu.current_stream`` (same function object). + On the first AOT/profile run, Dynamo builds its in-graph handler table and + registers ``torch.cuda.current_stream`` and ``torch.xpu.current_stream`` + separately; duplicate identity triggers:: + + AssertionError: Handler already registered for + + That surfaced as EngineCore failing in ``profile_run`` / ``_get_handlers()``. + The fix uses distinct shim callables so both can be registered. + + This test replays the post-init state (wrapper applied, patches left on + ``torch.cuda``) and checks that Dynamo's real ``_get_handlers()`` succeeds. + """ + # Same entry point as XPUModelRunner.__init__ (patches persist after exit). + with _torch_cuda_wrapper(): + pass + + # Fresh handler table build, as on first torch.compile / AOT in the worker. + # Registers torch.cuda.current_stream and torch.xpu.current_stream separately; + # if they are the same object (pre-fix alias), raises Handler already registered. + TorchInGraphFunctionVariable._get_handlers.cache_clear() + TorchInGraphFunctionVariable._get_handlers() diff --git a/tools/pre_commit/check_torch_cuda.py b/tools/pre_commit/check_torch_cuda.py index 9a67a013f1b..aec7b85d59c 100644 --- a/tools/pre_commit/check_torch_cuda.py +++ b/tools/pre_commit/check_torch_cuda.py @@ -9,7 +9,7 @@ import regex as re # --------------------------------------------------------------------------- # _TORCH_CUDA_PATTERNS = [ r"\btorch\.cuda\.(empty_cache|synchronize|device_count|current_device|memory_reserved|memory_allocated|max_memory_allocated|max_memory_reserved|reset_peak_memory_stats|memory_stats|mem_get_info|set_device|device\()\b", - r"\btorch\.cuda\.(manual_seed|manual_seed_all|Event)\b", + r"\btorch\.cuda\.(manual_seed|manual_seed_all)\b", r"\bwith\storch\.cuda\.device\b", # Calls torch.cuda.{_is_compiled/_device_count_amdsmi/_device_count_nvml} internally r"\bcuda_device_count_stateless\(\)\b", @@ -21,7 +21,6 @@ ALLOWED_FILES = { "vllm/device_allocator/", "vllm/distributed/weight_transfer/ipc_engine.py", "tests/distributed/test_packed_tensor.py", - "tools/pre_commit/check_torch_cuda.py", } @@ -40,13 +39,6 @@ def scan_file(path: str) -> int: f"Found {matched_text} API call. Use set_random_seed instead." ) return 1 - if matched_text == "torch.cuda.Event": - print( - f"{path}:{line_num}: " - "\033[91merror:\033[0m " - "Found torch.cuda.Event API call. Use torch.Event instead." - ) - return 1 print( f"{path}:{line_num}: " "\033[91merror:\033[0m " # red color diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 40ccbcf9a28..bed4d8254a0 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -2,12 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import functools from collections.abc import Callable -from contextlib import contextmanager -from typing import Protocol import torch from torch._ops import OpOverload -from torch.distributed import ProcessGroup import vllm.envs as envs from vllm.platforms import current_platform @@ -52,42 +49,6 @@ def is_aiter_found() -> bool: IS_AITER_FOUND = is_aiter_found() -class AiterCustomAllreduceProto(Protocol): - max_size: int - world_size: int - fully_connected: bool - - @contextmanager - def capture(self): ... - def close(self) -> None: ... - def fused_ar_rms( - self, - inp: torch.Tensor, - res_inp: torch.Tensor, - *, - w: torch.Tensor, - eps: float, - registered: bool = False, - use_1stage: bool = False, - ) -> tuple[torch.Tensor, torch.Tensor]: ... - def fused_ar_rms_per_group_quant( - self, - inp: torch.Tensor, - res_inp: torch.Tensor, - *, - w: torch.Tensor, - eps: float, - group_size: int = 128, - registered: bool = False, - use_1stage: bool = False, - emit_bf16: bool = False, - ) -> ( - tuple[torch.Tensor, torch.Tensor, torch.Tensor] - | tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] - ): ... - def should_custom_ar(self, inp: torch.Tensor) -> bool: ... - - def is_aiter_found_and_supported() -> bool: """Check if AITER library is available and platform supports it. @@ -830,6 +791,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl( ) -> tuple[torch.Tensor, torch.Tensor]: aiter_ar = rocm_aiter_ops.get_aiter_allreduce() assert aiter_ar is not None, "aiter allreduce must be initialized" + ca = aiter_ar.aiter_ca total_bytes = input_.numel() * input_.element_size() hidden_dim = input_.shape[-1] @@ -840,8 +802,8 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl( else: hidden_ok = False token_ok = token_num <= 80 - world_size = aiter_ar.world_size - full_nvlink = aiter_ar.fully_connected + world_size = ca.world_size + full_nvlink = ca.fully_connected if world_size == 2: size_ok = True @@ -854,12 +816,11 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl( use_1stage = hidden_ok and token_ok and size_ok - result = aiter_ar.fused_ar_rms( + result = ca.custom_fused_ar_rms( input_, residual, - w=weight, - eps=epsilon, - registered=torch.cuda.is_current_stream_capturing(), + weight, + epsilon, use_1stage=use_1stage, ) assert result is not None @@ -890,6 +851,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl( """ aiter_ar = rocm_aiter_ops.get_aiter_allreduce() assert aiter_ar is not None, "aiter allreduce must be initialized" + ca = aiter_ar.aiter_ca total_bytes = input_.numel() * input_.element_size() hidden_dim = input_.shape[-1] @@ -900,8 +862,8 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl( else: hidden_ok = False token_ok = token_num <= 80 - world_size = aiter_ar.world_size - full_nvlink = aiter_ar.fully_connected + world_size = ca.world_size + full_nvlink = ca.fully_connected if world_size == 2: size_ok = True @@ -914,7 +876,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl( use_1stage = hidden_ok and token_ok and size_ok - result = aiter_ar.fused_ar_rms_per_group_quant( + result = ca.fused_ar_rms_per_group_quant( input_, residual, w=weight, @@ -962,6 +924,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_impl( """ aiter_ar = rocm_aiter_ops.get_aiter_allreduce() assert aiter_ar is not None, "aiter allreduce must be initialized" + ca = aiter_ar.aiter_ca total_bytes = input_.numel() * input_.element_size() hidden_dim = input_.shape[-1] @@ -972,8 +935,8 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_impl( else: hidden_ok = False token_ok = token_num <= 80 - world_size = aiter_ar.world_size - full_nvlink = aiter_ar.fully_connected + world_size = ca.world_size + full_nvlink = ca.fully_connected if world_size == 2: size_ok = True @@ -986,7 +949,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_impl( use_1stage = hidden_ok and token_ok and size_ok - result = aiter_ar.fused_ar_rms_per_group_quant( + result = ca.fused_ar_rms_per_group_quant( input_, residual, w=weight, @@ -1577,6 +1540,7 @@ class rocm_aiter_ops: # Check if the env variable is set _AITER_ENABLED = envs.VLLM_ROCM_USE_AITER + _CUSTOM_ALL_REDUCE_ENABLED = envs.VLLM_ROCM_USE_AITER_CUSTOM_AR _LINEAR_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR _FMOE_ENABLED = envs.VLLM_ROCM_USE_AITER_MOE _MLA_ENABLED = envs.VLLM_ROCM_USE_AITER_MLA @@ -1598,9 +1562,6 @@ class rocm_aiter_ops: # num_shared_experts / shared_expert_scoring_func args (7-arg form). _TOPK_SOFTMAX_FUSED_SIGMOID: bool | None = None - _ALL_REDUCE_MAX_SIZE: int = 8192 * 1024 * 8 * 2 - _CUSTOM_ALL_REDUCE: AiterCustomAllreduceProto | None = None - @classmethod def refresh_env_variables(cls): """ @@ -1611,6 +1572,7 @@ class rocm_aiter_ops: you can call this function to reload the env variables. """ cls._AITER_ENABLED = envs.VLLM_ROCM_USE_AITER + cls._CUSTOM_ALL_REDUCE_ENABLED = envs.VLLM_ROCM_USE_AITER_CUSTOM_AR cls._LINEAR_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR cls._FMOE_ENABLED = envs.VLLM_ROCM_USE_AITER_MOE cls._MLA_ENABLED = envs.VLLM_ROCM_USE_AITER_MLA @@ -1770,6 +1732,11 @@ class rocm_aiter_ops: def is_mha_enabled(cls) -> bool: return cls._AITER_ENABLED and cls._MHA_ENABLED + @classmethod + @if_aiter_supported + def is_custom_all_reduce_enabled(cls) -> bool: + return cls._AITER_ENABLED and cls._CUSTOM_ALL_REDUCE_ENABLED + @classmethod @if_aiter_supported def is_shuffle_kv_cache_enabled(cls) -> bool: @@ -1824,33 +1791,20 @@ class rocm_aiter_ops: return cls.is_linear_enabled() and on_gfx950() @classmethod - def initialize_aiter_allreduce( - cls, group: ProcessGroup, device: torch.device - ) -> None: - try: - from aiter.dist.device_communicators.custom_all_reduce import ( - CustomAllreduce as AiterCustomAllreduce, - ) + def get_aiter_allreduce(cls): + """Return the TP device communicator's AITER custom-allreduce if it has + one, return None otherwise + """ + from vllm.distributed.device_communicators.aiter_custom_all_reduce import ( + AiterCustomAllreduce, + ) + from vllm.distributed.parallel_state import get_tp_group - cls._CUSTOM_ALL_REDUCE = AiterCustomAllreduce(group, device) - except Exception: - cls._CUSTOM_ALL_REDUCE = None - - @classmethod - def get_aiter_allreduce(cls) -> AiterCustomAllreduceProto | None: - return cls._CUSTOM_ALL_REDUCE - - @classmethod - def destroy_aiter_allreduce(cls) -> None: - if cls._CUSTOM_ALL_REDUCE is not None: - cls._CUSTOM_ALL_REDUCE.close() - cls._CUSTOM_ALL_REDUCE = None - - @classmethod - def get_aiter_allreduce_max_size(cls) -> int | None: - # effective max input size (based on upstream aiter version: v0.1.10.post3) - # https://github.com/ROCm/aiter/blob/6a0e7b26ccf33164785531212cc2ec2cde0b9243/aiter/dist/device_communicators/custom_all_reduce.py#L272-L273 - return int(cls._ALL_REDUCE_MAX_SIZE / 2) + device_comm = get_tp_group().device_communicator + aiter_ar_comm = getattr(device_comm, "aiter_ar_comm", None) + return ( + aiter_ar_comm if isinstance(aiter_ar_comm, AiterCustomAllreduce) else None + ) @classmethod @if_aiter_supported @@ -2165,21 +2119,6 @@ class rocm_aiter_ops: def get_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_op() -> OpOverload: # noqa: E501 return torch.ops.vllm.rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm.default # noqa: E501 - # TODO(frida-andersson): drop once vLLM pins AITER >= 0.1.14 (ROCm/aiter#2823). - @classmethod - def has_fused_allreduce_rmsnorm_quant_per_group(cls) -> bool: - """True if the running AITER build exposes the per-group AR+RMS+quant - kernel (added in ROCm/aiter PR #2823). - - The pattern registration in ``RocmAiterAllReduceFusionPass`` keys off - this so vLLM degrades to the AR+RMS-only fusion when run against an - older aiter that lacks the per-group launcher. - """ - aiter_ar = cls.get_aiter_allreduce() - return aiter_ar is not None and hasattr( - aiter_ar, "fused_ar_rms_per_group_quant" - ) - @staticmethod def get_fused_mla_dual_rms_norm_op() -> OpOverload: return torch.ops.vllm.fused_mla_dual_rms_norm.default diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 2a32a38a695..129e198423c 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2974,6 +2974,24 @@ if hasattr(torch.ops._C, "fused_experts_cpu"): return torch.empty_like(hidden_states) +if hasattr(torch.ops._C, "dynamic_4bit_int_moe"): + + @register_fake("_C::dynamic_4bit_int_moe") + def dynamic_4bit_int_moe_fake( + x: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + w13_packed: torch.Tensor, + w2_packed: torch.Tensor, + hidden_size: int, + intermediate_size: int, + group_size: int, + apply_router_weight_on_input: bool, + activation_kind: int, + ) -> torch.Tensor: + return x.new_empty((x.size(0), hidden_size)) + + def fused_experts_cpu( hidden_states: torch.Tensor, w1: torch.Tensor, diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index ee706037abb..ab400028925 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py @@ -19,7 +19,6 @@ from vllm.compilation.passes.fusion.rms_quant_fusion import ( from vllm.config import VllmConfig from vllm.config.utils import Range from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce -from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce from vllm.distributed.parallel_state import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, @@ -129,6 +128,29 @@ _FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB: dict[int, dict[int, float]] = { }, } +MiB = 1024 * 1024 + + +def _select_flashinfer_allreduce_use_oneshot( + workspace_backend: str, + device_capability: int | None, + world_size: int, + current_tensor_size: int, +) -> bool | None: + if workspace_backend == "mnnvl": + # FlashInfer sizes MNNVL workspaces around its own AUTO strategy. + # Forcing vLLM's per-rank threshold can request one-shot for tensors + # larger than the MNNVL one-shot workspace. + return None + + if device_capability is None: + max_one_shot_size = None + else: + max_one_shot_size = _FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB.get( + device_capability, {} + ).get(world_size) + return max_one_shot_size is None or current_tensor_size <= max_one_shot_size * MiB + if flashinfer_comm is not None: from vllm.distributed.device_communicators.flashinfer_all_reduce import ( @@ -139,8 +161,6 @@ if flashinfer_comm is not None: ar_fusion_patterns = flashinfer_comm.AllReduceFusionPattern - MiB = 1024 * 1024 - def call_trtllm_fused_allreduce_norm( allreduce_in: torch.Tensor, residual: torch.Tensor, @@ -175,16 +195,6 @@ if flashinfer_comm is not None: ) curr_device = current_platform.get_device_capability() device_capability = curr_device.to_int() if curr_device is not None else None - # Get one shot input size limit for the current world size - # for the current device capability - max_one_shot_size = _FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB.get( - device_capability, # type: ignore[arg-type, unused-ignore] - {}, - ).get(world_size, None) - # Use one shot if no max size is specified - use_oneshot = ( - max_one_shot_size is None or current_tensor_size <= max_one_shot_size * MiB - ) # Select workspace based on pattern: quant patterns use the # trtllm quant workspace, non-quant patterns use the primary workspace. @@ -206,6 +216,12 @@ if flashinfer_comm is not None: assert workspace is not None, ( "Flashinfer allreduce workspace must be initialized when using flashinfer" ) + use_oneshot = _select_flashinfer_allreduce_use_oneshot( + workspace.backend, + device_capability, + world_size, + current_tensor_size, + ) assert flashinfer_comm is not None if norm_out is None: norm_out = allreduce_in @@ -249,7 +265,7 @@ if flashinfer_comm is not None: # the end for the one-shot path; the two-shot path is synchronized # and keeps the early completion. Related one-shot instability in # the same kernel: flashinfer-ai/flashinfer#1223. - trigger_completion_at_end=use_oneshot + trigger_completion_at_end=(use_oneshot is True) or num_tokens > PDL_ADVANCE_LAUNCH_TOKENS, ) @@ -1473,39 +1489,23 @@ class RocmAiterAllReduceFusionPass(VllmFusionPatternMatcherPass): ) return - device_comm = get_tp_group().device_communicator - if device_comm is None: - logger.warning_once("Device communicator is required.") - return - - ca_comm = getattr(device_comm, "ca_comm", None) + ca_comm = rocm_aiter_ops.get_aiter_allreduce() if ca_comm is None: - logger.warning_once("Custom Allreduce is required.") + logger.warning_once( + "AITER allreduce fusions are disabled " + "because AITER Custom All Reduce is not enabled. " + "Set VLLM_ROCM_USE_AITER_CUSTOM_AR=1 " + "to enable it." + ) return self.ca_comm = ca_comm - assert isinstance(ca_comm, CustomAllreduce) - - group = get_tp_group().cpu_group - rocm_aiter_ops.initialize_aiter_allreduce(group, self.device) hidden_dim = config.model_config.get_hidden_size() element_size = torch.tensor([], dtype=self.model_dtype).element_size() - max_size = rocm_aiter_ops.get_aiter_allreduce_max_size() - if max_size is None: - logger.warning("AITER allreduce fusion must be initialized") - return - - # Aiter's fused_allreduce_rmsnorm kernel dispatches on hidden_dim. - # Before aiter v0.1.12 the launcher was template-specialized on HIDDEN_DIM - # and silently no-op'd for sizes outside {512, 1024, 2048, 4096}. From v0.1.12 - # hidden_dim is a runtime argument. Detect the older API via the missing - # `_pool` attribute and skip fusion for unsupported sizes. - # Ref (old kernel): https://github.com/ROCm/aiter/blob/6a0e7b26ccf33164785531212cc2ec2cde0b9243/csrc/include/custom_all_reduce.cuh#L2590 - aiter_ar = rocm_aiter_ops.get_aiter_allreduce() + max_size = ca_comm.effective_max_size() _AITER_OLD_FUSED_AR_RMS_HIDDEN = (512, 1024, 2048, 4096) if ( - aiter_ar is not None - and not hasattr(aiter_ar, "_pool") + not ca_comm.supports_dynamic_hidden_dim and hidden_dim not in _AITER_OLD_FUSED_AR_RMS_HIDDEN ): logger.warning_once( @@ -1515,10 +1515,6 @@ class RocmAiterAllReduceFusionPass(VllmFusionPatternMatcherPass): _AITER_OLD_FUSED_AR_RMS_HIDDEN, hidden_dim, ) - # Tear down aiter's custom-allreduce so its IPC handles don't - # race with vllm's ca_comm on the unfused fallback path. - with contextlib.suppress(Exception): - rocm_aiter_ops.destroy_aiter_allreduce() return max_token_num = max_size // (hidden_dim * element_size) @@ -1532,9 +1528,7 @@ class RocmAiterAllReduceFusionPass(VllmFusionPatternMatcherPass): # fall back to the AR+RMS-only fusion paired with PR #41825's # standalone RMS+quant fusion -- still correct, just leaves the # post-AR quant as a standalone kernel. - supports_per_group_quant = ( - rocm_aiter_ops.has_fused_allreduce_rmsnorm_quant_per_group() - ) + supports_per_group_quant = ca_comm.supports_per_group_quant if not supports_per_group_quant: logger.warning_once( "AITER AR+RMS+per-group-FP8-quant fusion disabled: aiter " @@ -1609,9 +1603,3 @@ class RocmAiterAllReduceFusionPass(VllmFusionPatternMatcherPass): logger.debug( "%s Replaced %s patterns", self.__class__.__name__, self.matched_count ) - - def __del__(self) -> None: - if getattr(self, "disabled", True): - return - with contextlib.suppress(Exception): - rocm_aiter_ops.destroy_aiter_allreduce() diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index e98ab1b3e08..aff94fe03d2 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -147,6 +147,7 @@ LinearBackend = Literal[ "flashinfer_cudnn", "flashinfer_b12x", "marlin", + "humming", "triton", "deep_gemm", "torch", diff --git a/vllm/config/model.py b/vllm/config/model.py index 2ec754f9346..6e3fef0dcca 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -74,6 +74,12 @@ else: logger = init_logger(__name__) +# Process-local record of which (arch, target) model-class overrides have been +# registered in *this* process. Must not live on ModelConfig: that instance is +# pickled to each worker, so an instance flag would arrive already "registered" +# while the worker's own global ModelRegistry is still untouched. +_REGISTERED_MODEL_CLASS_OVERRIDES: set[tuple[str, str]] = set() + RunnerOption = Literal["auto", RunnerType] ConvertType = Literal["none", "embed", "classify"] ConvertOption = Literal["auto", ConvertType] @@ -274,6 +280,13 @@ class ModelConfig: hf_overrides: HfOverrides = field(default_factory=dict) """If a dictionary, contains arguments to be forwarded to the Hugging Face config. If a callable, it is called to update the HuggingFace config.""" + model_class_overrides: dict[str, str] = field(default_factory=dict) + """Override the model class used for one or more architectures, mapping the + architecture name to a `"module:class"` target (the same format accepted by + `ModelRegistry.register_model`). This registers the target class at runtime, + e.g. `{"GlmMoeDsaForCausalLM": + "vllm.models.deepseek_v32.nvidia.model:DeepseekV32ForCausalLM"}`. This + argument is for development and debugging purposes only.""" generation_config: str = "auto" """The folder path to the generation config. Defaults to `"auto"`, the generation config will be loaded from model path. If set to `"vllm"`, no @@ -812,8 +825,34 @@ class ModelConfig: @property def registry(self): + self._maybe_register_model_class_overrides() return me_models.ModelRegistry + def _maybe_register_model_class_overrides(self) -> None: + # Apply ``model_class_overrides`` here because this property is the + # single chokepoint through which every model-class inspect/resolve + # passes, in both the engine front-end and every worker process. The + # guard is process-local (see ``_REGISTERED_MODEL_CLASS_OVERRIDES``), so + # each worker re-registers into its own ModelRegistry exactly once + # rather than trusting a pickled-in instance flag. + if not self.model_class_overrides: + return + pending = [ + (arch, target) + for arch, target in self.model_class_overrides.items() + if (arch, target) not in _REGISTERED_MODEL_CLASS_OVERRIDES + ] + if not pending: + return + logger.warning_once( + "Applying model_class_overrides %s. This is intended for " + "development/debugging.", + str(self.model_class_overrides), + ) + for arch, target in pending: + me_models.ModelRegistry.register_model(arch, target) + _REGISTERED_MODEL_CLASS_OVERRIDES.add((arch, target)) + @property def architectures(self) -> list[str]: return self.model_arch_config.architectures @@ -999,6 +1038,7 @@ class ModelConfig: "modelopt", "modelopt_fp4", "modelopt_mxfp8", + "mxfp8", "modelopt_mixed", # Ensure heavy backends are probed last to avoid unnecessary # imports during override detection (e.g., MXFP4 imports Triton) diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 983759bc9a5..2e0126368cf 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import copy +import functools +from collections.abc import Callable from typing import TYPE_CHECKING, Any, Literal, get_args from pydantic import Field, SkipValidation, field_validator, model_validator @@ -46,6 +48,7 @@ MTPModelTypes = Literal[ "qwen3_5_mtp", "longcat_flash_mtp", "minimax_m3_mtp", + "bailing_hybrid_mtp", "mtp", "pangu_ultra_moe_mtp", "step3p5_mtp", @@ -463,6 +466,21 @@ class SpeculativeConfig: {"n_predict": n_predict, "architectures": ["Qwen3NextMTP"]} ) + architectures = getattr(hf_config, "architectures", []) or [] + if ( + hf_config.model_type == "bailing_hybrid" + or "BailingMoeV2_5ForCausalLM" in architectures + ): + hf_config.model_type = "bailing_hybrid_mtp" + if hf_config.model_type == "bailing_hybrid_mtp": + n_predict = getattr(hf_config, "num_nextn_predict_layers", None) + hf_config.update( + { + "n_predict": n_predict, + "architectures": ["BailingMoeV25MTPModel"], + } + ) + if hf_config.model_type == "exaone_moe": hf_config.model_type = "exaone_moe_mtp" if hf_config.model_type == "exaone_moe_mtp": @@ -572,6 +590,40 @@ class SpeculativeConfig: return hf_config + @staticmethod + def _apply_composed_hf_override( + target_hf_overrides: Callable[[PretrainedConfig], PretrainedConfig], + hf_config: PretrainedConfig, + ) -> PretrainedConfig: + hf_config = SpeculativeConfig.hf_config_override(hf_config) + return target_hf_overrides(hf_config) + + @staticmethod + def compose_draft_hf_overrides( + target_hf_overrides: HfOverrides | None, + ) -> Callable[[PretrainedConfig], PretrainedConfig]: + """Build the ``hf_overrides`` for the draft ``ModelConfig``. + + Callable overrides on the target are config-to-config transforms + (e.g. test harnesses shrinking ``num_hidden_layers``) and must also + reach the draft config — otherwise a draft belonging to a large + target is instantiated at full size even when the target is shrunk. + Dict overrides are target-specific key patches and are not applied + to the draft. + + The composed override must stay picklable: the draft ``ModelConfig`` + is sent to spawned engine-core processes, so a local closure would + fail with ``Can't get local object`` during pickling. Bind the + target via ``functools.partial`` over a module-referenceable static + method instead. + """ + if not callable(target_hf_overrides): + return SpeculativeConfig.hf_config_override + + return functools.partial( + SpeculativeConfig._apply_composed_hf_override, target_hf_overrides + ) + def __post_init__(self): # Note: "method" is a new parameter that helps to extend the # configuration of non-model-based proposers, and the "model" parameter @@ -736,7 +788,12 @@ class SpeculativeConfig: if self.method == "medusa": draft_hf_overrides = {"model_type": "medusa"} else: - draft_hf_overrides = SpeculativeConfig.hf_config_override + # Compose any callable hf_overrides set on the target so the + # draft config receives the same transform (e.g. the test + # shrink). Dict overrides stay target-only. + draft_hf_overrides = SpeculativeConfig.compose_draft_hf_overrides( + self.target_model_config.hf_overrides + ) self.draft_model_config = ModelConfig( model=self.model, runner="draft", diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 8f90bb23326..e8796a3db8e 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1850,8 +1850,13 @@ class VllmConfig: tp_size = self.parallel_config.tensor_parallel_size from vllm._aiter_ops import rocm_aiter_ops - if rocm_aiter_ops.is_enabled(): - max_size = rocm_aiter_ops.get_aiter_allreduce_max_size() + max_size: int | None = None + if rocm_aiter_ops.is_custom_all_reduce_enabled(): + from vllm.distributed.device_communicators.aiter_custom_all_reduce import ( # noqa: E501 + AiterCustomAllreduce, + ) + + max_size = AiterCustomAllreduce.effective_max_size() else: max_size = compilation_config.pass_config.flashinfer_max_size(tp_size) if max_size is not None and self.model_config is not None: diff --git a/vllm/distributed/device_communicators/aiter_custom_all_reduce.py b/vllm/distributed/device_communicators/aiter_custom_all_reduce.py new file mode 100644 index 00000000000..63e06b77a77 --- /dev/null +++ b/vllm/distributed/device_communicators/aiter_custom_all_reduce.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""vLLM-owned wrapper over AITER's ``CustomAllreduce``. + +vLLM's ``CudaCommunicator`` stores one of these as ``aiter_ar_comm`` (when +``VLLM_ROCM_USE_AITER_CUSTOM_AR`` is set) so the plain allreduce and +the fused allreduce+RMSNorm path share a single AITER instance with its IPC buffers. + +""" + +import torch +from torch.distributed import ProcessGroup + +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +class AiterCustomAllreduce: + # Default IPC buffer size for AITER's CustomAllreduce. + MAX_SIZE: int = 8192 * 1024 * 8 * 2 + + @classmethod + def effective_max_size(cls) -> int: + """ + Max input byte size eligible for AITER custom allreduce. + """ + return cls.MAX_SIZE // 2 + + def __init__( + self, + group: ProcessGroup, + device: int | str | torch.device, + max_size: int | None = None, + ): + from aiter.dist.device_communicators.custom_all_reduce import ( + CustomAllreduce as _AiterCustomAllreduce, + ) + + if max_size is None: + max_size = self.MAX_SIZE + + self._impl = _AiterCustomAllreduce(group, device, max_size=max_size) + + @property + def aiter_ca(self): + return self._impl + + @property + def disabled(self) -> bool: + return self._impl.disabled + + def should_custom_ar(self, inp: torch.Tensor) -> bool: + return self._impl.should_custom_ar(inp) + + def custom_all_reduce(self, inp: torch.Tensor) -> torch.Tensor | None: + return self._impl.custom_all_reduce(inp) + + def capture(self): + return self._impl.capture() + + def close(self) -> None: + self._impl.close() + + @property + def supports_dynamic_hidden_dim(self) -> bool: + """Aiter's fused_allreduce_rmsnorm kernel dispatches on hidden_dim. + Before aiter v0.1.12 the launcher was template-specialized on HIDDEN_DIM + and silently no-op'd for sizes outside {512, 1024, 2048, 4096}. From v0.1.12 + hidden_dim is a runtime argument. Older builds are detected via + AiterCustomAllreduce.supports_dynamic_hidden_dim; This function is used to + skip fusion for unsupported sizes on them. + Ref (old kernel): https://github.com/ROCm/aiter/blob/6a0e7b26ccf33164785531212cc2ec2cde0b9243/csrc/include/custom_all_reduce.cuh#L2590 + """ + return hasattr(self._impl, "_pool") + + @staticmethod + def build_supports_per_group_quant() -> bool: + """True if the running AITER build exposes the per-group AR+RMS+quant + kernel (added in ROCm/aiter PR #2823). + + The pattern registration in ``RocmAiterAllReduceFusionPass`` keys off + this so vLLM degrades to the AR+RMS-only fusion when run against an + older aiter that lacks the per-group launcher. + """ + from aiter.dist.device_communicators.custom_all_reduce import ( + CustomAllreduce as _AiterCustomAllreduce, + ) + + return hasattr(_AiterCustomAllreduce, "fused_ar_rms_per_group_quant") + + # TODO(frida-andersson): drop once vLLM pins AITER >= 0.1.14 (ROCm/aiter#2823). + @property + def supports_per_group_quant(self) -> bool: + return self.build_supports_per_group_quant() diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index b92015b1880..555fd0ec948 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -6,6 +6,7 @@ import torch from torch.distributed import ProcessGroup import vllm.envs as envs +from vllm._aiter_ops import rocm_aiter_ops from vllm.distributed.device_communicators.all_reduce_utils import ( NCCL_SYMM_MEM_ALL_REDUCE_CONFIG, should_nccl_symm_mem_ag_rs, @@ -19,6 +20,7 @@ from vllm.logger import init_logger from vllm.platforms import current_platform from ..utils import StatelessProcessGroup +from .aiter_custom_all_reduce import AiterCustomAllreduce from .base_device_communicator import DeviceCommunicatorBase logger = init_logger(__name__) @@ -48,16 +50,21 @@ class CudaCommunicator(DeviceCommunicatorBase): use_custom_allreduce = False use_torch_symm_mem = False use_flashinfer_allreduce = False + use_aiter_allreduce = False else: from vllm.distributed.parallel_state import _ENABLE_CUSTOM_ALL_REDUCE use_custom_allreduce = _ENABLE_CUSTOM_ALL_REDUCE use_torch_symm_mem = envs.VLLM_ALLREDUCE_USE_SYMM_MEM use_flashinfer_allreduce = envs.VLLM_ALLREDUCE_USE_FLASHINFER + use_aiter_allreduce = use_custom_allreduce and bool( + rocm_aiter_ops.is_custom_all_reduce_enabled() + ) self.use_custom_allreduce = use_custom_allreduce self.use_torch_symm_mem = use_torch_symm_mem self.use_flashinfer_allreduce = use_flashinfer_allreduce + self.use_aiter_allreduce = use_aiter_allreduce # lazy import to avoid documentation build error from vllm.distributed.device_communicators.custom_all_reduce import ( @@ -85,6 +92,7 @@ class CudaCommunicator(DeviceCommunicatorBase): self.qr_comm: QuickAllReduce | None = None self.symm_mem_comm: SymmMemCommunicator | None = None self.fi_ar_comm: FlashInferAllReduce | None = None + self.aiter_ar_comm: AiterCustomAllreduce | None = None if use_torch_symm_mem and current_platform.is_cuda(): self.symm_mem_comm = SymmMemCommunicator( @@ -98,7 +106,13 @@ class CudaCommunicator(DeviceCommunicatorBase): device=self.device, ) - if use_custom_allreduce and self.world_size > 1: + if self.use_aiter_allreduce and self.world_size > 1: + self.aiter_ar_comm = AiterCustomAllreduce( + group=self.cpu_group, + device=self.device, + ) + + if use_custom_allreduce and self.aiter_ar_comm is None and self.world_size > 1: # Initialize a custom fast all-reduce implementation. self.ca_comm = CustomAllreduce( group=self.cpu_group, @@ -108,13 +122,14 @@ class CudaCommunicator(DeviceCommunicatorBase): ), ) - if current_platform.is_rocm(): - # Initialize a custom quick all-reduce implementation for AMD. - # Quick reduce is designed as a complement to custom allreduce. - # Based on quickreduce (https://github.com/mk1-project/quickreduce). - # If it's a rocm, 'use_custom_allreduce==True' means it must - # currently be an MI300 series. - self.qr_comm = QuickAllReduce(group=self.cpu_group, device=self.device) + if use_custom_allreduce and self.world_size > 1 and current_platform.is_rocm(): + # Initialize a custom quick all-reduce implementation for AMD. + # Quick reduce is designed as a complement to custom allreduce + # (vLLM's or AITER's), so it is initialized for either backend. + # Based on quickreduce (https://github.com/mk1-project/quickreduce). + # On ROCm, 'use_custom_allreduce==True' means it must currently be + # an MI300 series. + self.qr_comm = QuickAllReduce(group=self.cpu_group, device=self.device) if self.world_size > 1: self._log_all_reduce_backend_selection() @@ -203,6 +218,7 @@ class CudaCommunicator(DeviceCommunicatorBase): "NCCL_SYMM_MEM", "QUICK_REDUCE", "FLASHINFER", + "AITER_CUSTOM", "CUSTOM", "SYMM_MEM", "PYNCCL", @@ -236,6 +252,8 @@ class CudaCommunicator(DeviceCommunicatorBase): enabled_ar_backends.append("QUICK_REDUCE") if self.fi_ar_comm is not None and not self.fi_ar_comm.disabled: enabled_ar_backends.append("FLASHINFER") + if self.aiter_ar_comm is not None and not self.aiter_ar_comm.disabled: + enabled_ar_backends.append("AITER_CUSTOM") if self.ca_comm is not None and not self.ca_comm.disabled: enabled_ar_backends.append("CUSTOM") if self.symm_mem_comm is not None and not self.symm_mem_comm.disabled: @@ -261,8 +279,8 @@ class CudaCommunicator(DeviceCommunicatorBase): out = torch.ops.vllm.all_reduce_symmetric_with_copy(input_) if out is not None: return out - # always try quick reduce first, then flashinfer, then custom allreduce, - # and then pynccl. (quick reduce just for ROCM MI3*) + # always try quick reduce first, then flashinfer, then the AITER or vLLM + # custom allreduce, and then pynccl. (quick reduce just for ROCM MI3*) qr_comm = self.qr_comm if ( qr_comm is not None @@ -281,6 +299,15 @@ class CudaCommunicator(DeviceCommunicatorBase): out = fi_ar_comm.all_reduce(input_) assert out is not None return out + aiter_ar_comm = self.aiter_ar_comm + if ( + aiter_ar_comm is not None + and not aiter_ar_comm.disabled + and aiter_ar_comm.should_custom_ar(input_) + ): + out = aiter_ar_comm.custom_all_reduce(input_) + assert out is not None + return out ca_comm = self.ca_comm if ( ca_comm is not None @@ -509,6 +536,9 @@ class CudaCommunicator(DeviceCommunicatorBase): self.pynccl_comm = None if self.ca_comm is not None: self.ca_comm = None + if self.aiter_ar_comm is not None: + self.aiter_ar_comm.close() + self.aiter_ar_comm = None if self.fi_ar_comm is not None: self.fi_ar_comm.destroy() self.fi_ar_comm = None diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index e98c765f537..feacb03d28b 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -742,8 +742,8 @@ class EplbState: is_main_rank = ep_rank == 0 if is_main_rank: if not self.is_async or is_profile: - start_event = torch.Event(enable_timing=True) - end_event = torch.Event(enable_timing=True) + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) start_event.record() logger.info( "Rearranging experts %s %s...", diff --git a/vllm/distributed/eplb/eplb_utils.py b/vllm/distributed/eplb/eplb_utils.py index 21a7ee68fa9..dee19749745 100644 --- a/vllm/distributed/eplb/eplb_utils.py +++ b/vllm/distributed/eplb/eplb_utils.py @@ -31,7 +31,7 @@ class CpuGpuEvent: """ def __init__(self): - self._event = torch.Event() + self._event = torch.cuda.Event() self._recorded = threading.Event() def wait(self, stream: torch.cuda.Stream | None = None): @@ -56,7 +56,7 @@ class CpuGpuEvent: "CpuGpuEvent.record() called before the previous event was " "consumed by wait()" ) - self._event = torch.Event() + self._event = torch.cuda.Event() self._event.record(stream) self._recorded.set() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py index a604bd5528f..299ff037ad2 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py @@ -239,7 +239,7 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): # this event is complete the request is considered "done sending" # by get_finished; clients block on the per-file flock to wait for # the disk write itself. - self._req_copy_events: dict[str, torch.Event] = {} + self._req_copy_events: dict[str, torch.cuda.Event] = {} # req_ids reported as finished-generating by the scheduler, # accumulated across get_finished calls. self._accumulated_finished_req_ids: set[str] = set() @@ -320,7 +320,7 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): @staticmethod def _write_tensors( tensors: dict[str, torch.Tensor], - event: torch.Event, + event: torch.cuda.Event, filename: str, lock_fd: int | None, ) -> None: @@ -375,7 +375,7 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): copy_stream = self._get_copy_stream() # Ensure the copy stream sees all prior writes on the default stream. - ready_event = torch.Event() + ready_event = torch.cuda.Event() ready_event.record() copy_stream.wait_event(ready_event) @@ -396,7 +396,7 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): pinned_hs.copy_(hidden_states_gpu, non_blocking=True) # Record completion of this copy on the copy stream. - copy_done = torch.Event() + copy_done = torch.cuda.Event() copy_done.record(copy_stream) # token_ids is already on CPU (created in request_finished). diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_client.py b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_client.py index edcf83b1925..a54233453bb 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_client.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_client.py @@ -221,7 +221,7 @@ class Hf3fsClient: @wsynchronized() def batch_write( - self, offsets: list[int], tensors: list[torch.Tensor], event: torch.Event + self, offsets: list[int], tensors: list[torch.Tensor], event: torch.cuda.Event ) -> list[int]: """Write data from tensors to the file at specified offsets. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py index 55a8b5a161e..526375952fe 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py @@ -133,7 +133,7 @@ class AsyncOperationManager: # CUDA streams for async operations self._save_stream = torch.cuda.Stream() self._load_stream = torch.cuda.Stream() - self._save_event = torch.Event() + self._save_event = torch.cuda.Event() # Buffer allocators for data copying self._save_buffer_allocator = CopyBufferAllocator( @@ -171,7 +171,7 @@ class AsyncOperationManager: def submit_save_operation(self, request_id: str, block_ids, block_hashes) -> Future: """Submit a save operation for async execution.""" future: Future[Any] = Future() - main_stream_event = torch.Event() + main_stream_event = torch.cuda.Event() main_stream_event.record() task = (request_id, block_ids, block_hashes, future, main_stream_event) self._save_queue.put(task) @@ -304,7 +304,7 @@ class AsyncOperationManager: block_ids, buffers, "gather" ) - save_stream_event = torch.Event() + save_stream_event = torch.cuda.Event() save_stream_event.record(self._save_stream) # Record gather completion # Step3: Write data in batches diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/hf3fs_mock_client.py b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/hf3fs_mock_client.py index e2718d1faa1..3914663a62d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/hf3fs_mock_client.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/hf3fs_mock_client.py @@ -75,7 +75,7 @@ class Hf3fsClient: return torch.frombuffer(buffer_data, dtype=dtype) def batch_write( - self, offsets: list[int], tensors: list[torch.Tensor], event: torch.Event + self, offsets: list[int], tensors: list[torch.Tensor], event: torch.cuda.Event ) -> list[int]: """Write data from tensors to file at specified offsets.""" results = [] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py index 6d83380cad3..2e75519df12 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py @@ -430,7 +430,7 @@ class LMCacheMPWorkerAdapter: @_lmcache_nvtx_annotate def submit_store_request( - self, request_id: str, op: LoadStoreOp, event: torch.Event + self, request_id: str, op: LoadStoreOp, event: torch.cuda.Event ): """ Submit a KV cache store request to LMCache @@ -464,7 +464,7 @@ class LMCacheMPWorkerAdapter: @_lmcache_nvtx_annotate def submit_retrieve_request( - self, request_id: str, op: LoadStoreOp, event: torch.Event + self, request_id: str, op: LoadStoreOp, event: torch.cuda.Event ): """ Submit a KV cache retrieve request to LMCache @@ -501,7 +501,7 @@ class LMCacheMPWorkerAdapter: self, request_ids: list[str], ops: list[LoadStoreOp], - event: torch.Event, + event: torch.cuda.Event, ): """ Submit a batched store request to LMCache @@ -550,7 +550,7 @@ class LMCacheMPWorkerAdapter: self, request_ids: list[str], ops: list[LoadStoreOp], - event: torch.Event, + event: torch.cuda.Event, ): """ Submit a batched retrieve request to LMCache diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py index 2ca35be2b51..8786e91a5a1 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py @@ -589,7 +589,7 @@ class LMCacheMPConnectorUpstream(KVConnectorBase_V1): return with torch.cuda.stream(torch.cuda.current_stream()): - event = torch.Event(interprocess=True) + event = torch.cuda.Event(interprocess=True) event.record() self.worker_adapter.batched_submit_retrieve_requests( @@ -663,7 +663,7 @@ class LMCacheMPConnectorUpstream(KVConnectorBase_V1): return with torch.cuda.stream(torch.cuda.current_stream()): - event = torch.Event(interprocess=True) + event = torch.cuda.Event(interprocess=True) event.record() self.worker_adapter.batched_submit_store_requests( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py index 3b69ce9a177..ef98ec0d4e4 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py @@ -323,7 +323,7 @@ class ReqMeta: can_save: bool | None = None load_spec: LoadSpec | None = None is_last_chunk: bool | None = None - current_event: torch.Event | None = None + current_event: torch.cuda.Event | None = None token_ids: list[int] | None = None num_prompt_tokens: int | None = None diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 818c2479a14..e60d2f47a4e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -1357,7 +1357,7 @@ class MooncakeStoreWorker: current_event = None for request in meta.requests: if request.can_save: - current_event = torch.Event() + current_event = torch.cuda.Event() current_event.record() break diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py index 7b8ab566058..15585123e5c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py @@ -56,7 +56,7 @@ class WriteTask: local_block_ids: list[int] remote_block_ids_hint: list[int] | None layer_name: str - event: torch.Event + event: torch.cuda.Event remote_notify_port: int remote_ip: str enqueue_time: float = field(default_factory=time.perf_counter) 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 172836f5cc3..de21a1398e0 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 @@ -1061,7 +1061,7 @@ class MoRIIOConnectorWorker: # when mori-io supports ibgda functionality stream = torch.cuda.current_stream() - event = torch.Event() + event = torch.cuda.Event() event.record(stream) task = WriteTask( diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 11b9e24e864..162ed03d23b 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -269,11 +269,17 @@ def _create_subgroups_split_group( must enter with the same ``split_ranks`` definition. Each rank receives the subgroup it belongs to. """ + from vllm.distributed.utils import ( + get_cpu_distributed_timeout_or_none, + get_distributed_timeout_or_none, + ) + device_backend_str = _device_backend_str(torch_distributed_backend) self_device_group = torch.distributed.split_group( split_ranks=group_ranks, group_desc=f"{group_name}:device", backend=device_backend_str, + timeout=get_distributed_timeout_or_none(), ) # CPU subgroup: split_group requires the requested backend filter to # include the parent's default device type (= the device the parent PG @@ -284,6 +290,7 @@ def _create_subgroups_split_group( split_ranks=group_ranks, group_desc=f"{group_name}:cpu", backend=f"cpu:gloo,{device_backend_str}", + timeout=get_cpu_distributed_timeout_or_none(), ) return self_device_group, self_cpu_group @@ -417,13 +424,19 @@ class GroupCoordinator: self.rank_in_group = ranks.index(self.rank) break else: - from vllm.distributed.utils import get_cpu_distributed_timeout_or_none + from vllm.distributed.utils import ( + get_cpu_distributed_timeout_or_none, + get_distributed_timeout_or_none, + ) timeout = get_cpu_distributed_timeout_or_none() + device_timeout = get_distributed_timeout_or_none() for ranks in group_ranks: device_group = torch.distributed.new_group( - ranks, backend=torch_distributed_backend + ranks, + backend=torch_distributed_backend, + timeout=device_timeout, ) # a group with `gloo` backend, to allow direct coordination between # processes through the CPU. @@ -504,10 +517,16 @@ class GroupCoordinator: This is a collective call: every world rank must invoke it. Used where we want to issue ops that can run concurrently with ops on `device_group`. """ + from vllm.distributed.utils import get_distributed_timeout_or_none + + device_timeout = get_distributed_timeout_or_none() sibling: ProcessGroup | None = None for ranks in self.group_ranks: pg = torch.distributed.new_group( - ranks, backend=self.torch_distributed_backend, group_desc=group_desc + ranks, + backend=self.torch_distributed_backend, + group_desc=group_desc, + timeout=device_timeout, ) if self.rank in ranks: sibling = pg diff --git a/vllm/distributed/utils.py b/vllm/distributed/utils.py index ef3c11ff64e..eec9890f028 100644 --- a/vllm/distributed/utils.py +++ b/vllm/distributed/utils.py @@ -533,6 +533,16 @@ def get_cpu_distributed_timeout_or_none() -> timedelta | None: return timedelta(seconds=timeout_seconds) if timeout_seconds is not None else None +def get_distributed_timeout_or_none() -> timedelta | None: + from vllm.config import get_current_vllm_config_or_none + + vllm_config = get_current_vllm_config_or_none() + if vllm_config is None: + return None + timeout_seconds = vllm_config.parallel_config.distributed_timeout_seconds + return timedelta(seconds=timeout_seconds) if timeout_seconds is not None else None + + def init_gloo_process_group( prefix_store: PrefixStore, group_rank: int, @@ -616,6 +626,10 @@ def stateless_init_torch_distributed_process_group( gloo_timeout = get_cpu_distributed_timeout_or_none() if gloo_timeout is not None: timeout = gloo_timeout + else: + device_timeout = get_distributed_timeout_or_none() + if device_timeout is not None: + timeout = device_timeout if listen_socket is not None: store = create_tcp_store( diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index efdd7696fdc..c7a9335bbeb 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -536,6 +536,9 @@ class EngineArgs: code_revision: str | None = ModelConfig.code_revision hf_token: bool | str | None = ModelConfig.hf_token hf_overrides: HfOverrides = get_field(ModelConfig, "hf_overrides") + model_class_overrides: dict[str, str] = get_field( + ModelConfig, "model_class_overrides" + ) tokenizer_revision: str | None = ModelConfig.tokenizer_revision quantization: QuantizationMethods | str | None = ModelConfig.quantization quantization_config: "dict[str, Any] | QuantizationConfigArgs | None" = None @@ -851,6 +854,9 @@ class EngineArgs: model_group.add_argument("--config-format", **model_kwargs["config_format"]) model_group.add_argument("--hf-token", **model_kwargs["hf_token"]) model_group.add_argument("--hf-overrides", **model_kwargs["hf_overrides"]) + model_group.add_argument( + "--model-class-overrides", **model_kwargs["model_class_overrides"] + ) model_group.add_argument("--pooler-config", **model_kwargs["pooler_config"]) model_group.add_argument( "--generation-config", **model_kwargs["generation_config"] @@ -1622,6 +1628,7 @@ class EngineArgs: code_revision=self.code_revision, hf_token=self.hf_token, hf_overrides=self.hf_overrides, + model_class_overrides=self.model_class_overrides, tokenizer_revision=self.tokenizer_revision, max_model_len=self.max_model_len, quantization=self.quantization, diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index 5a26e475b06..062a7947e79 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -130,6 +130,7 @@ async def init_generate_state( enable_force_include_usage=args.enable_force_include_usage, enable_log_outputs=args.enable_log_outputs, enable_log_deltas=args.enable_log_deltas, + enable_per_request_metrics=args.enable_per_request_metrics, ) state.openai_serving_chat = ( OpenAIServingChat(**_chat_kwargs) if "generate" in supported_tasks else None @@ -150,6 +151,7 @@ async def init_generate_state( return_tokens_as_token_ids=args.return_tokens_as_token_ids, enable_prompt_tokens_details=args.enable_prompt_tokens_details, enable_force_include_usage=args.enable_force_include_usage, + enable_per_request_metrics=args.enable_per_request_metrics, ) if "generate" in supported_tasks else None diff --git a/vllm/entrypoints/generate/base/serving.py b/vllm/entrypoints/generate/base/serving.py index 196b6e3f87f..78f84ca872b 100644 --- a/vllm/entrypoints/generate/base/serving.py +++ b/vllm/entrypoints/generate/base/serving.py @@ -18,6 +18,7 @@ from vllm.entrypoints.openai.completion.protocol import CompletionRequest from vllm.entrypoints.openai.engine.protocol import ( ErrorResponse, GenerationError, + PerRequestTimingMetrics, ) from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.openai.responses.protocol import ResponsesRequest @@ -34,6 +35,7 @@ from vllm.tracing import ( extract_trace_headers, log_tracing_disabled_warning, ) +from vllm.v1.metrics.stats import RequestStateStats logger = init_logger(__name__) @@ -41,6 +43,61 @@ RequestT = TypeVar("RequestT", bound=AnyRequest) _T = TypeVar("_T") +def build_per_request_timing_metrics( + metrics: RequestStateStats | None, + num_generation_tokens: int, +) -> PerRequestTimingMetrics: + """Build per-request timing metrics from ``RequestStateStats``. + + ``generation_time_ms`` is the decode interval only (first output token to + last output token); it excludes both queue wait and prefill/TTFT. + ``tokens_per_second`` is overall output throughput: all generated tokens + over the inference interval (scheduling to last output token), so it counts + the prefill/TTFT phase and is not simply the reciprocal of ``mean_itl_ms``. + Each field is left ``None`` when the timestamps it depends on are + unavailable. + """ + if metrics is None: + return PerRequestTimingMetrics() + + queued_ts = metrics.queued_ts + scheduled_ts = metrics.scheduled_ts + first_token_ts = metrics.first_token_ts + last_token_ts = metrics.last_token_ts + + time_to_first_token_ms: float | None = None + generation_time_ms: float | None = None + queue_time_ms: float | None = None + mean_itl_ms: float | None = None + tokens_per_second: float | None = None + + if scheduled_ts > 0 and first_token_ts > 0: + time_to_first_token_ms = (first_token_ts - scheduled_ts) * 1000 + + if first_token_ts > 0 and last_token_ts > 0: + generation_time_ms = (last_token_ts - first_token_ts) * 1000 + + if queued_ts > 0 and scheduled_ts > 0: + queue_time_ms = (scheduled_ts - queued_ts) * 1000 + + if first_token_ts > 0 and last_token_ts > 0 and num_generation_tokens > 1: + decode_time = last_token_ts - first_token_ts + mean_itl_ms = decode_time / (num_generation_tokens - 1) * 1000 + + if scheduled_ts > 0 and last_token_ts > 0: + inference_time_ms = (last_token_ts - scheduled_ts) * 1000 + if inference_time_ms > 0: + tokens_per_second = num_generation_tokens / inference_time_ms * 1000 + + return PerRequestTimingMetrics( + time_to_first_token_ms=time_to_first_token_ms, + generation_time_ms=generation_time_ms, + queue_time_ms=queue_time_ms, + mean_itl_ms=mean_itl_ms, + tokens_per_second=tokens_per_second, + ) + + @dataclass(kw_only=True) class ServeContext(Generic[RequestT]): request: RequestT diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py index 6fb27c365d9..71dc508a34a 100644 --- a/vllm/entrypoints/openai/api_server.py +++ b/vllm/entrypoints/openai/api_server.py @@ -51,7 +51,7 @@ from vllm.entrypoints.serve.utils.server_utils import ( log_response, validation_exception_handler, ) -from vllm.exceptions import VLLMValidationError +from vllm.exceptions import VLLMUnprocessableEntityError, VLLMValidationError from vllm.logger import init_logger from vllm.reasoning import ReasoningParserManager from vllm.renderers.online_derenderer import OnlineDerenderer @@ -245,6 +245,7 @@ def build_app( app.exception_handler(EngineDeadError)(engine_error_handler) app.exception_handler(GenerationError)(generation_error_handler) app.exception_handler(VLLMValidationError)(exception_handler) + app.exception_handler(VLLMUnprocessableEntityError)(exception_handler) app.exception_handler(Exception)(exception_handler) # Ensure --api-key option from CLI takes precedence over VLLM_API_KEY diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 22cae80ce84..cce51157f84 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -26,6 +26,7 @@ from vllm.entrypoints.openai.engine.protocol import ( FunctionDefinition, LegacyStructuralTagResponseFormat, OpenAIBaseModel, + PerRequestTimingMetrics, StreamOptions, StructuralTagResponseFormat, ToolCall, @@ -133,6 +134,7 @@ class ChatCompletionResponse(OpenAIBaseModel): kv_transfer_params: dict[str, Any] | None = Field( default=None, description="KVTransfer parameters." ) + metrics: PerRequestTimingMetrics | None = None class ChatCompletionResponseStreamChoice(OpenAIBaseModel): @@ -160,6 +162,7 @@ class ChatCompletionStreamResponse(OpenAIBaseModel): # Rendered prompt text from chat templating (only set when # ``return_prompt_text=True`` on the request); only sent on the first chunk. prompt_text: str | None = None + metrics: PerRequestTimingMetrics | None = None class ChatCompletionToolsParam(OpenAIBaseModel): diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 9ef144e6a14..caa3f724da4 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -22,6 +22,7 @@ from vllm.entrypoints.chat_utils import ( from vllm.entrypoints.generate.base.serving import ( GenerateBaseServing, GenerationError, + build_per_request_timing_metrics, clamp_prompt_logprobs, format_token_id_placeholder, ) @@ -41,6 +42,7 @@ from vllm.entrypoints.openai.engine.protocol import ( DeltaMessage, ErrorResponse, FunctionCall, + PerRequestTimingMetrics, PromptTokenUsageInfo, RequestResponseMetadata, ToolCall, @@ -123,6 +125,7 @@ class OpenAIServingChat(GenerateBaseServing): enable_log_outputs: bool = False, enable_log_deltas: bool = True, default_chat_template_kwargs: dict[str, Any] | None = None, + enable_per_request_metrics: bool = False, ) -> None: super().__init__( engine_client=engine_client, @@ -161,6 +164,7 @@ class OpenAIServingChat(GenerateBaseServing): self.enable_prompt_tokens_details = enable_prompt_tokens_details self.enable_force_include_usage = enable_force_include_usage + self.enable_per_request_metrics = enable_per_request_metrics self.default_sampling_params = self.model_config.get_diff_sampling_param() mc = self.model_config self.override_max_tokens = ( @@ -461,8 +465,10 @@ class OpenAIServingChat(GenerateBaseServing): stream_options, self.enable_force_include_usage ) + last_res: RequestOutput | None = None try: async for res in result_generator: + last_res = res if res.prompt_token_ids is not None: num_prompt_tokens = len(res.prompt_token_ids) if res.encoder_prompt_token_ids is not None: @@ -752,6 +758,21 @@ class OpenAIServingChat(GenerateBaseServing): mm_token_counts, ) + # In streaming, metrics ride on this final usage chunk, which is + # only emitted when usage reporting is enabled (i.e. + # ``stream_options.include_usage=true`` or + # ``--enable-force-include-usage``). + stream_per_request_metrics: PerRequestTimingMetrics | None = None + if ( + self.enable_per_request_metrics + # See note in chat_completion_full_generator: suppress for n>1. + and (request.n or 1) == 1 + ): + last_metrics = last_res.metrics if last_res is not None else None + stream_per_request_metrics = build_per_request_timing_metrics( + last_metrics, completion_tokens + ) + final_usage_chunk = ChatCompletionStreamResponse( id=request_id, object=chunk_object_type, @@ -760,6 +781,7 @@ class OpenAIServingChat(GenerateBaseServing): model=model_name, usage=final_usage, system_fingerprint=self.system_fingerprint, + metrics=stream_per_request_metrics, ) final_usage_data = final_usage_chunk.model_dump_json( exclude_unset=True, exclude_none=True @@ -1003,6 +1025,18 @@ class OpenAIServingChat(GenerateBaseServing): request_metadata.final_usage_info = usage + per_request_metrics: PerRequestTimingMetrics | None = None + if ( + self.enable_per_request_metrics + # Timing metrics describe a single generation stream. For n>1 the + # returned stats belong to only one of the n sequences, so they + # cannot be accurately attributed to the request; suppress instead. + and (request.n or 1) == 1 + ): + per_request_metrics = build_per_request_timing_metrics( + final_res.metrics, num_generated_tokens + ) + # ``final_res.prompt`` is the rendered chat-templated prompt text prompt_text = final_res.prompt if request.return_prompt_text else None @@ -1019,6 +1053,7 @@ class OpenAIServingChat(GenerateBaseServing): ), prompt_text=prompt_text, kv_transfer_params=final_res.kv_transfer_params, + metrics=per_request_metrics, ) # Log complete response if output logging is enabled diff --git a/vllm/entrypoints/openai/cli_args.py b/vllm/entrypoints/openai/cli_args.py index 1533895edcd..8dbb6994390 100644 --- a/vllm/entrypoints/openai/cli_args.py +++ b/vllm/entrypoints/openai/cli_args.py @@ -25,12 +25,9 @@ from vllm.entrypoints.serve.utils.constants import ( H11_MAX_HEADER_COUNT_DEFAULT, H11_MAX_INCOMPLETE_EVENT_SIZE_DEFAULT, ) -from vllm.logger import init_logger from vllm.tool_parsers import ToolParserManager from vllm.utils.argparse_utils import FlexibleArgumentParser -logger = init_logger(__name__) - class LoRAParserAction(argparse.Action): def __call__( @@ -134,6 +131,8 @@ class BaseFrontendArgs: log. The default of None means unlimited.""" enable_prompt_tokens_details: bool = False """If set to True, enable prompt_tokens_details in usage.""" + enable_per_request_metrics: bool = False + """If set to True, include per-request timing metrics in API responses.""" enable_server_load_tracking: bool = False """If set to True, enable tracking server_load_metrics in the app state.""" enable_force_include_usage: bool = False @@ -398,6 +397,14 @@ def validate_parsed_serve_args(args: argparse.Namespace): if args.enable_log_outputs and not args.enable_log_requests: raise TypeError("Error: --enable-log-outputs requires --enable-log-requests") + if getattr(args, "enable_per_request_metrics", False) and getattr( + args, "disable_log_stats", False + ): + raise ValueError( + "Error: --enable-per-request-metrics requires engine statistics " + "logging; remove --disable-log-stats to enable per-request metrics." + ) + if args.data_parallel_multi_port_external_lb: from vllm.entrypoints.openai.dp_supervisor import ( validate_multi_port_external_lb_args, diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index b96d4f3c0c7..23bf4dd83f3 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -15,6 +15,7 @@ from vllm.entrypoints.openai.engine.protocol import ( AnyResponseFormat, LegacyStructuralTagResponseFormat, OpenAIBaseModel, + PerRequestTimingMetrics, StreamOptions, StructuralTagResponseFormat, UsageInfo, @@ -558,6 +559,7 @@ class CompletionResponse(OpenAIBaseModel): kv_transfer_params: dict[str, Any] | None = Field( default=None, description="KVTransfer parameters." ) + metrics: PerRequestTimingMetrics | None = None class CompletionResponseStreamChoice(OpenAIBaseModel): @@ -589,3 +591,4 @@ class CompletionStreamResponse(OpenAIBaseModel): # Set only on the final chunk of a stream to mirror non-streaming responses # without the per-chunk serialization overhead. system_fingerprint: str | None = None + metrics: PerRequestTimingMetrics | None = None diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index aeade306465..d26a455cc8e 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -16,6 +16,7 @@ from vllm.engine.protocol import EngineClient from vllm.entrypoints.generate.base.serving import ( GenerateBaseServing, GenerationError, + build_per_request_timing_metrics, clamp_prompt_logprobs, format_token_id_placeholder, ) @@ -29,6 +30,7 @@ from vllm.entrypoints.openai.completion.protocol import ( ) from vllm.entrypoints.openai.engine.protocol import ( ErrorResponse, + PerRequestTimingMetrics, PromptTokenUsageInfo, RequestResponseMetadata, UsageInfo, @@ -61,6 +63,7 @@ class OpenAIServingCompletion(GenerateBaseServing): return_tokens_as_token_ids: bool = False, enable_prompt_tokens_details: bool = False, enable_force_include_usage: bool = False, + enable_per_request_metrics: bool = False, ): super().__init__( engine_client=engine_client, @@ -72,6 +75,7 @@ class OpenAIServingCompletion(GenerateBaseServing): self.online_renderer = online_renderer self.enable_prompt_tokens_details = enable_prompt_tokens_details self.enable_force_include_usage = enable_force_include_usage + self.enable_per_request_metrics = enable_per_request_metrics self.default_sampling_params = self.model_config.get_diff_sampling_param() mc = self.model_config @@ -300,8 +304,10 @@ class OpenAIServingCompletion(GenerateBaseServing): stream_options, self.enable_force_include_usage ) + last_res: RequestOutput | None = None try: async for prompt_idx, res in result_generator: + last_res = res prompt_token_ids = res.prompt_token_ids prompt_logprobs = res.prompt_logprobs @@ -448,6 +454,23 @@ class OpenAIServingCompletion(GenerateBaseServing): ) if include_usage: + # In streaming, metrics ride on this final usage chunk, which is + # only emitted when usage reporting is enabled (i.e. + # ``stream_options.include_usage=true`` or + # ``--enable-force-include-usage``). + stream_per_request_metrics: PerRequestTimingMetrics | None = None + if ( + self.enable_per_request_metrics + # See note in request_output_to_completion_response: suppress + # when not attributable to one stream (multi-prompt or n>1). + and num_prompts == 1 + and (request.n or 1) == 1 + ): + last_metrics = last_res.metrics if last_res is not None else None + stream_per_request_metrics = build_per_request_timing_metrics( + last_metrics, total_completion_tokens + ) + final_usage_chunk = CompletionStreamResponse( id=request_id, created=created_time, @@ -455,6 +478,7 @@ class OpenAIServingCompletion(GenerateBaseServing): choices=[], usage=final_usage_info, system_fingerprint=self.system_fingerprint, + metrics=stream_per_request_metrics, ) final_usage_data = final_usage_chunk.model_dump_json( exclude_unset=False, exclude_none=True @@ -589,6 +613,23 @@ class OpenAIServingCompletion(GenerateBaseServing): ) request_metadata.final_usage_info = usage + + per_request_metrics: PerRequestTimingMetrics | None = None + if ( + self.enable_per_request_metrics + # Metrics describe a single generation stream, so suppress them when + # they cannot be attributed to one: multiple prompts (timestamps + # span prompts) or n>1 (stats belong to one of the n sequences). + and len(final_res_batch) == 1 + and (request.n or 1) == 1 + ): + last_metrics = ( + last_final_res.metrics if last_final_res is not None else None + ) + per_request_metrics = build_per_request_timing_metrics( + last_metrics, num_generated_tokens + ) + if final_res_batch: kv_transfer_params = final_res_batch[0].kv_transfer_params return CompletionResponse( @@ -599,6 +640,7 @@ class OpenAIServingCompletion(GenerateBaseServing): usage=usage, system_fingerprint=self.system_fingerprint, kv_transfer_params=kv_transfer_params, + metrics=per_request_metrics, ) def _create_completion_logprobs( diff --git a/vllm/entrypoints/openai/engine/protocol.py b/vllm/entrypoints/openai/engine/protocol.py index 084d8d429a6..2c32fcf20c6 100644 --- a/vllm/entrypoints/openai/engine/protocol.py +++ b/vllm/entrypoints/openai/engine/protocol.py @@ -115,6 +115,14 @@ class UsageInfo(OpenAIBaseModel): prompt_tokens_details: PromptTokenUsageInfo | None = None +class PerRequestTimingMetrics(OpenAIBaseModel): + time_to_first_token_ms: float | None = None + generation_time_ms: float | None = None + queue_time_ms: float | None = None + mean_itl_ms: float | None = None + tokens_per_second: float | None = None + + class RequestResponseMetadata(BaseModel): request_id: str final_usage_info: UsageInfo | None = None diff --git a/vllm/entrypoints/openai/responses/context.py b/vllm/entrypoints/openai/responses/context.py index 92aad86b211..ed989e2ba9f 100644 --- a/vllm/entrypoints/openai/responses/context.py +++ b/vllm/entrypoints/openai/responses/context.py @@ -349,6 +349,7 @@ class ParsableContext(ConversationContext): reasoning=reasoning, content=content, tool_calls=tool_calls, + tools=self.request.tools, ) ) elif completion.text: diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index eb2d66bdd8f..ba8bc5a40f1 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -592,10 +592,19 @@ class ResponsesRequest(OpenAIBaseModel): ) elif is_named_tool_choice and tools is not None: tool_name = tool_choice.get("name") - tool_names = { - t.get("name") if isinstance(t, dict) else getattr(t, "name", None) - for t in tools - } + tool_names = set() + for tool in tools: + if isinstance(tool, dict): + if tool.get("type") == "namespace": + namespace = tool.get("name") + for namespaced_tool in tool.get("tools", []): + namespaced_name = namespaced_tool.get("name") + tool_names.add(namespaced_name) + tool_names.add(f"{namespace}__{namespaced_name}") + else: + tool_names.add(tool.get("name")) + else: + tool_names.add(getattr(tool, "name", None)) if not tool_name or tool_name not in tool_names: raise VLLMValidationError( "Tool choice 'function' not found in 'tools' parameter.", diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index a62c4623992..40a52012792 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -1073,6 +1073,7 @@ class OpenAIServingResponses(GenerateBaseServing): content=content, tool_calls=tool_calls, logprobs=logprobs, + tools=request.tools, ) # Fallback when no parser is configured @@ -1339,7 +1340,7 @@ class OpenAIServingResponses(GenerateBaseServing): [StreamingResponsesResponse], StreamingResponsesResponse ], ) -> AsyncGenerator[StreamingResponsesResponse, None]: - processor = SimpleStreamingEventProcessor() + processor = SimpleStreamingEventProcessor(tools=request.tools) def _get_logprobs( output: CompletionOutput, diff --git a/vllm/entrypoints/openai/responses/streaming_events.py b/vllm/entrypoints/openai/responses/streaming_events.py index 35021caf2ab..531a35c5722 100644 --- a/vllm/entrypoints/openai/responses/streaming_events.py +++ b/vllm/entrypoints/openai/responses/streaming_events.py @@ -58,6 +58,7 @@ from openai.types.responses.response_output_item import McpCall from openai.types.responses.response_reasoning_item import ( Content as ResponseReasoningTextContent, ) +from openai.types.responses.tool import Tool from openai_harmony import Message as HarmonyMessage from vllm.entrypoints.mcp.tool_server import ToolServer @@ -71,6 +72,10 @@ from vllm.entrypoints.openai.responses.protocol import ( ResponseReasoningPartDoneEvent, StreamingResponsesResponse, ) +from vllm.entrypoints.openai.responses.utils import ( + build_responses_tool_call_name_map, + resolve_responses_tool_call_name, +) from vllm.outputs import CompletionOutput from vllm.parser.harmony import Segment from vllm.utils import random_uuid @@ -815,6 +820,7 @@ class SimpleStreamingState: accumulated_text: str = "" tool_call_id: str = "" tool_call_name: str = "" + tool_call_namespace: str | None = None tool_call_index: int | None = None has_emitted_tool_call_delta: bool = False current_state: _StateType = field(default_factory=lambda: _StateType.NONE) @@ -1016,11 +1022,13 @@ def emit_simple_tool_call_open( state: SimpleStreamingState, name: str, index: int | None, + namespace: str | None = None, ) -> list[StreamingResponsesResponse]: state.current_state = _StateType.TOOL_CALL state.current_item_id = random_uuid() state.tool_call_id = f"call_{random_uuid()}" state.tool_call_name = name + state.tool_call_namespace = namespace state.tool_call_index = index state.accumulated_text = "" state.has_emitted_tool_call_delta = False @@ -1034,6 +1042,7 @@ def emit_simple_tool_call_open( id=state.current_item_id, call_id=state.tool_call_id, name=name, + namespace=namespace, arguments="", status="in_progress", ), @@ -1081,6 +1090,7 @@ def emit_simple_tool_call_done( item=ResponseFunctionToolCall( type="function_call", name=state.tool_call_name, + namespace=state.tool_call_namespace, arguments=state.accumulated_text, status="completed", id=state.current_item_id, @@ -1089,6 +1099,7 @@ def emit_simple_tool_call_done( ), ) state.output_index += 1 + state.tool_call_namespace = None state.current_state = _StateType.NONE return events @@ -1166,8 +1177,13 @@ class SimpleStreamingEventProcessor: ), } - def __init__(self, state: SimpleStreamingState | None = None) -> None: + def __init__( + self, + state: SimpleStreamingState | None = None, + tools: list[Tool] | None = None, + ) -> None: self.state = state or SimpleStreamingState() + self.tool_call_name_map = build_responses_tool_call_name_map(tools) def resolve_target_state( self, delta_message: DeltaMessage @@ -1224,8 +1240,15 @@ class SimpleStreamingEventProcessor: handlers = self._STATE_HANDLERS[target_state] if target_state == _StateType.TOOL_CALL: assert tool_call is not None + call_name = resolve_responses_tool_call_name( + tool_call.function.name, + tool_call_name_map=self.tool_call_name_map, + ) return handlers.open_fn( - self.state, tool_call.function.name, tool_call.index + self.state, + call_name.name, + tool_call.index, + call_name.namespace, ) return handlers.open_fn(self.state) diff --git a/vllm/entrypoints/openai/responses/utils.py b/vllm/entrypoints/openai/responses/utils.py index a2f35dca235..07a9704f9b5 100644 --- a/vllm/entrypoints/openai/responses/utils.py +++ b/vllm/entrypoints/openai/responses/utils.py @@ -35,6 +35,12 @@ from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionMessa from vllm.entrypoints.openai.engine.protocol import FunctionCall from vllm.entrypoints.openai.responses.protocol import ResponseInputOutputItem from vllm.logger import init_logger +from vllm.tool_parsers.utils import ( + build_responses_tool_call_name_map, + flat_namespace_tool_name, + iter_response_function_tool_dicts, + resolve_responses_tool_call_name, +) from vllm.utils import random_uuid logger = init_logger(__name__) @@ -45,8 +51,10 @@ def build_response_output_items( content: str | None, tool_calls: list[FunctionCall] | None, logprobs: list[Logprob] | None = None, + tools: list[Tool] | None = None, ) -> list[ResponseOutputItem]: outputs: list[ResponseOutputItem] = [] + tool_call_name_map = build_responses_tool_call_name_map(tools) if reasoning: outputs.append( @@ -81,6 +89,9 @@ def build_response_output_items( if tool_calls: for idx, tool_call in enumerate(tool_calls): + call_name = resolve_responses_tool_call_name( + tool_call.name, tool_call_name_map=tool_call_name_map + ) outputs.append( ResponseFunctionToolCall( id=f"fc_{random_uuid()}", @@ -88,7 +99,8 @@ def build_response_output_items( or make_tool_call_id(func_name=tool_call.name, idx=idx), type="function_call", status="completed", - name=tool_call.name, + name=call_name.name, + namespace=call_name.namespace, arguments=tool_call.arguments, ) ) @@ -219,10 +231,13 @@ def _construct_message_from_response_item( ) if isinstance(item, ResponseFunctionToolCall): + tool_name = item.name + if item.namespace: + tool_name = flat_namespace_tool_name(item.namespace, item.name) tool_call = ChatCompletionMessageToolCallParam( id=item.call_id, function=FunctionCallTool( - name=item.name, + name=tool_name, arguments=item.arguments, ), type="function", @@ -318,7 +333,17 @@ def _construct_message_from_response_item( def extract_function_tool_names(tools: list[Tool]) -> frozenset[str]: - return frozenset(tool.name for tool in tools if tool.type == "function") + names = [] + for tool in tools: + if tool.type == "function": + names.append(tool.name) + elif tool.type == "namespace": + names.extend( + flat_namespace_tool_name(tool.name, namespaced_tool.name) + for namespaced_tool in tool.tools + if namespaced_tool.type == "function" + ) + return frozenset(names) def extract_tool_types(tools: list[Tool]) -> set[str]: @@ -358,7 +383,7 @@ def construct_tool_dicts( tool_dicts = None else: tool_dicts = [ - convert_tool_responses_to_completions_format(tool.model_dump()) - for tool in tools + convert_tool_responses_to_completions_format(tool) + for tool in iter_response_function_tool_dicts(tools) ] return tool_dicts diff --git a/vllm/entrypoints/openai/run_batch.py b/vllm/entrypoints/openai/run_batch.py index 58975b4f86b..6ae608da0ad 100644 --- a/vllm/entrypoints/openai/run_batch.py +++ b/vllm/entrypoints/openai/run_batch.py @@ -28,6 +28,7 @@ from urllib3.util import parse_url import vllm.envs as envs from vllm.config import config +from vllm.connections import global_http_connection from vllm.engine.arg_utils import AsyncEngineArgs from vllm.engine.protocol import EngineClient from vllm.entrypoints.openai.api_server import init_app_state @@ -493,18 +494,9 @@ async def download_bytes_from_url( # between urllib3 and aiohttp (e.g. backslash-@ attacks). url = url_spec.url - async with ( - aiohttp.ClientSession() as session, - session.get( - url, - allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS, - ) as resp, - ): - if resp.status != 200: - raise Exception( - f"Failed to download data from URL: {url}. Status: {resp.status}" - ) - return await resp.read() + return await global_http_connection.async_get_bytes( + url, allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS + ) else: raise ValueError( diff --git a/vllm/entrypoints/serve/utils/error_response.py b/vllm/entrypoints/serve/utils/error_response.py index 4dea1513a42..fc17a75c75a 100644 --- a/vllm/entrypoints/serve/utils/error_response.py +++ b/vllm/entrypoints/serve/utils/error_response.py @@ -27,12 +27,20 @@ def create_error_response( "create_error_response called with %s: %s", type(exc).__name__, exc ) - from vllm.exceptions import VLLMNotFoundError, VLLMValidationError + from vllm.exceptions import ( + VLLMNotFoundError, + VLLMUnprocessableEntityError, + VLLMValidationError, + ) if isinstance(exc, VLLMValidationError): err_type = "BadRequestError" status_code = HTTPStatus.BAD_REQUEST param = exc.parameter + elif isinstance(exc, VLLMUnprocessableEntityError): + err_type = "UnprocessableEntityError" + status_code = HTTPStatus.UNPROCESSABLE_ENTITY + param = exc.parameter elif isinstance(exc, VLLMNotFoundError): err_type = "NotFoundError" status_code = HTTPStatus.NOT_FOUND diff --git a/vllm/envs.py b/vllm/envs.py index 1f94be8ac5f..13a19b86c0a 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -119,6 +119,7 @@ if TYPE_CHECKING: VLLM_USE_OINK_OPS: bool = False VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD: bool = True VLLM_ROCM_USE_AITER: bool = False + VLLM_ROCM_USE_AITER_CUSTOM_AR: bool = True VLLM_ROCM_USE_AITER_PAGED_ATTN: bool = False VLLM_ROCM_USE_AITER_LINEAR: bool = True VLLM_ROCM_USE_AITER_LINEAR_HIPBMM: bool = False @@ -1146,6 +1147,11 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_ROCM_USE_AITER": lambda: ( os.getenv("VLLM_ROCM_USE_AITER", "False").lower() in ("true", "1") ), + # Use AITER's CustomAllreduce as the custom-allreduce backend inside vLLM's + # CudaCommunicator on ROCm. + "VLLM_ROCM_USE_AITER_CUSTOM_AR": lambda: ( + os.getenv("VLLM_ROCM_USE_AITER_CUSTOM_AR", "True").lower() in ("true", "1") + ), # Whether to use aiter paged attention. # By default is disabled. "VLLM_ROCM_USE_AITER_PAGED_ATTN": lambda: ( diff --git a/vllm/exceptions.py b/vllm/exceptions.py index 931040b8ceb..4112c3de24b 100644 --- a/vllm/exceptions.py +++ b/vllm/exceptions.py @@ -64,3 +64,37 @@ class LoRAAdapterNotFoundError(VLLMNotFoundError): def __str__(self): return self.message + + +class VLLMUnprocessableEntityError(ValueError): + """vLLM-specific error for unprocessable entity requests. + + This exception is raised when the request content is invalid or cannot be + processed, such as when an image URL points to a non-existent or inaccessible + resource (404, 403, DNS failure, etc.). + + Args: + message: The error message describing the unprocessable entity. + parameter: Optional parameter name that failed validation. + value: Optional value that was rejected during validation. + """ + + def __init__( + self, + message: str, + *, + parameter: str | None = None, + value: Any = None, + ) -> None: + super().__init__(message) + self.parameter = parameter + self.value = value + + def __str__(self): + base = super().__str__() + extras = [] + if self.parameter is not None: + extras.append(f"parameter={self.parameter}") + if self.value is not None: + extras.append(f"value={self.value}") + return f"{base} ({', '.join(extras)})" if extras else base diff --git a/vllm/kernels/helion/config_manager.py b/vllm/kernels/helion/config_manager.py index ca37a68e810..052b515c37f 100644 --- a/vllm/kernels/helion/config_manager.py +++ b/vllm/kernels/helion/config_manager.py @@ -162,13 +162,6 @@ class ConfigSet: config_key, ) - def has_config(self, platform: str, config_key: CaseKey) -> bool: - platform = platform.lower() - platform_dict = self._configs.get(platform) - if platform_dict is None: - return False - return config_key in platform_dict - class ConfigManager: """File-level configuration management for Helion kernels (global singleton).""" @@ -327,15 +320,3 @@ class ConfigManager: logger.info("Saved config to: %s", platform_path) return platform_path - - def config_exists( - self, - kernel_name: str, - platform: str, - config_key: CaseKey, - ) -> bool: - platform_data = self._load_platform_file(kernel_name, platform) - if not platform_data: - return False - target = dict(config_key) - return any(entry["key"] == target for entry in platform_data) diff --git a/vllm/kernels/helion/configs/silu_and_mul_per_block_quant/nvidia_b200.json b/vllm/kernels/helion/configs/silu_and_mul_per_block_quant/nvidia_b200.json new file mode 100644 index 00000000000..b50cabb0da9 --- /dev/null +++ b/vllm/kernels/helion/configs/silu_and_mul_per_block_quant/nvidia_b200.json @@ -0,0 +1,2099 @@ +[ + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 0, + 2 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 4 + ], + "range_warp_specializes": [ + null + ], + "range_multi_buffers": [ + false + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "", + "" + ], + "num_warps": 2, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "persistent_interleaved", + "num_sm_multiplier": 4, + "maxnreg": 128 + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "last", + "first" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "last" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "" + ], + "num_warps": 1, + "num_stages": 2, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "first" + ], + "num_warps": 1, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "", + "last" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "", + "last" + ], + "num_warps": 2, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "", + "last" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "", + "last" + ], + "num_warps": 1, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "last", + "last" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "", + "last" + ], + "num_warps": 2, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "", + "" + ], + "num_warps": 4, + "num_stages": 4, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "last", + "last" + ], + "num_warps": 2, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "last" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "last", + "last" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "last" + ], + "num_warps": 2, + "num_stages": 4, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "last", + "last" + ], + "num_warps": 2, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "last", + "last" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "last", + "last" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + } +] \ No newline at end of file diff --git a/vllm/kernels/helion/configs/silu_and_mul_per_block_quant/nvidia_h100.json b/vllm/kernels/helion/configs/silu_and_mul_per_block_quant/nvidia_h100.json new file mode 100644 index 00000000000..ca411c69470 --- /dev/null +++ b/vllm/kernels/helion/configs/silu_and_mul_per_block_quant/nvidia_h100.json @@ -0,0 +1,2207 @@ +[ + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "last" + ], + "num_warps": 4, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "last" + ], + "num_warps": 4, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "", + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "last" + ], + "num_warps": 4, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 8, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "", + "last" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "" + ], + "num_warps": 8, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "last" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "last" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "last" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "last", + "first" + ], + "num_warps": 4, + "num_stages": 7, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 0, + 2, + 1 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 7, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "", + "" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 2, + "num_stages": 4, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "first" + ], + "num_warps": 8, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "", + "" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "first" + ], + "num_warps": 8, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "", + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "", + "" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 16384 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "last" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 16384 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "", + "last" + ], + "num_warps": 1, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 16384 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "", + "" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + } +] \ No newline at end of file diff --git a/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py b/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py new file mode 100644 index 00000000000..f3aaf226b04 --- /dev/null +++ b/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from itertools import product +from typing import Any + +import torch + +from vllm.kernels.helion.case_key import CaseKey +from vllm.kernels.helion.utils import ( + get_fp8_dtype, + get_int8_min_max, + get_int8_min_scaling_factor, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + get_fp8_min_max, +) +from vllm.platforms import current_platform +from vllm.utils.import_utils import has_helion + +if not has_helion(): + raise ImportError( + "Helion kernel requires helion to be installed. " + "Install it with: pip install helion" + ) + +import helion +import helion.language as hl + +from vllm.kernels.helion.register import register_kernel + +logger = init_logger(__name__) + + +def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]: + # TODO(xiaohongchen1991): it is difficult for kernel author to cover all input + # property combination. Currently, dtypes are fixed. We need optimization to + # bucket/skip some combinations + num_tokens_list = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192] + intermediate_size_list = [6144, 12288, 25600] + + in_dtype: torch.dtype = torch.bfloat16 + out_dtype: torch.dtype = current_platform.fp8_dtype() + scale_dtype: torch.dtype = torch.float32 + group_size_list = [128] + inputs = {} + for intermediate_size, group_size, num_tokens in product( + intermediate_size_list, group_size_list, num_tokens_list + ): + input = torch.randn( + num_tokens, 2 * intermediate_size, device="cuda", dtype=in_dtype + ) + result = torch.empty( + num_tokens, intermediate_size, device=input.device, dtype=out_dtype + ) + scale = torch.empty( + (num_tokens, intermediate_size // group_size), + device=input.device, + dtype=scale_dtype, + ) + scale_ub = torch.mean(input).to(scale_dtype) + + config_key = CaseKey( + { + "intermediate_size": intermediate_size, + "group_size": group_size, + "num_tokens": num_tokens, + } + ) + inputs[config_key] = (result, input, scale, group_size, scale_ub, False) + + return inputs + + +_pick_cache: dict[tuple[int, int, int], CaseKey | None] = {} + + +def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey | None: + """Pick the best pre-tuned config for the given input shape. + + Selection strategy: + 1. Find the closest intermediate_size among available configs + (exact match preferred). + 2. Find the closest group_size among available configs + (exact match preferred). + 3. Among the num_tokens values tuned for that intermediate_size and group_size, + pick the smallest num_tokens >= the input's num_tokens. If the input is + larger than all available num_tokens, fall back to the largest. + """ + + if not config_keys: + return None + + result, _, _, group_size, *_ = args + num_tokens, intermediate_size = result.shape + + cache_key = (num_tokens, group_size, intermediate_size) + cached = _pick_cache.get(cache_key) + if cached is not None: + return cached + + configs: dict[int, dict[int, list[int]]] = {} + for key in config_keys: + if key.is_default(): + continue + configs.setdefault(key["intermediate_size"], {}).setdefault( + key["group_size"], [] + ).append(key["num_tokens"]) + + if not configs: + return None + + best_intermediate_size = min(configs, key=lambda s: abs(s - intermediate_size)) + best_group_size = min( + configs[best_intermediate_size], key=lambda s: abs(s - group_size) + ) + available_num_tokens = sorted(configs[best_intermediate_size][best_group_size]) + best_num_tokens = next( + (n for n in available_num_tokens if n >= num_tokens), available_num_tokens[-1] + ) + + result = CaseKey( + { + "intermediate_size": best_intermediate_size, + "group_size": best_group_size, + "num_tokens": best_num_tokens, + } + ) + _pick_cache[cache_key] = result + return result + + +def fake_impl( + out: torch.Tensor, # [num_tokens, intermediate_size] + input: torch.Tensor, # [num_tokens, 2 * intermediate_size] + scales: torch.Tensor, # [num_tokens, groups_per_row] + group_size: int, + scale_ub: torch.Tensor | None = None, # scalar tensor + is_scale_transposed: bool = False, +) -> None: + return + + +def baseline( + out: torch.Tensor, # [num_tokens, intermediate_size] + input: torch.Tensor, # [num_tokens, 2 * intermediate_size] + scales: torch.Tensor, # [num_tokens, groups_per_row] + group_size: int, + scale_ub: torch.Tensor | None = None, # scalar tensor + is_scale_transposed: bool = False, +) -> None: + torch.ops._C.silu_and_mul_per_block_quant( + out, input, scales, group_size, scale_ub, is_scale_transposed + ) + + +@register_kernel( + mutates_args=["out", "scales"], + config_picker=pick_config, + input_generator=generate_inputs, + fake_impl=fake_impl, + helion_settings=helion.Settings( + autotune_baseline_fn=baseline, + ignore_warnings=[helion.exc.TensorOperationInWrapper], + ), +) # type: ignore[misc] +def silu_and_mul_per_block_quant( + out: torch.Tensor, # [num_tokens, intermediate_size] + input: torch.Tensor, # [num_tokens, 2 * intermediate_size] + scales: torch.Tensor, # [num_tokens, groups_per_row] + group_size: int, + scale_ub: torch.Tensor | None = None, # scalar tensor + is_scale_transposed: bool = False, # dummy +) -> None: + # This code assumes batch_dim and num_tokens are flattened + assert input.ndim == 2 + num_tokens, two_intermediate_size = input.shape + hl.specialize(two_intermediate_size) + + assert two_intermediate_size % 2 == 0 + intermediate_size = two_intermediate_size // 2 + + assert out.shape[0] == num_tokens + assert out.shape[1] == intermediate_size + fp8_dtype = get_fp8_dtype() + assert out.dtype in [fp8_dtype, torch.int8] + + if scale_ub is not None: + assert out.dtype == fp8_dtype + assert scale_ub.dtype == torch.float32 + + assert scales.ndim == 2 and scales.dtype == torch.float32 + + assert scales.shape[0] == num_tokens + groups_per_row = scales.shape[1] + hl.specialize(groups_per_row) + assert ( + intermediate_size % group_size == 0 + and intermediate_size // group_size == groups_per_row + ) + + assert group_size in [64, 128] + hl.specialize(group_size) + + assert input.stride()[-1] == 1 + assert out.stride()[-1] == 1 + + quant_dtype = out.dtype + qtype_traits_min: int | float + qtype_traits_max: int | float + if quant_dtype == torch.int8: + qtype_traits_min, qtype_traits_max = get_int8_min_max() + min_scaling_factor = get_int8_min_scaling_factor() + else: + qtype_traits_min, qtype_traits_max = get_fp8_min_max() + min_scaling_factor = 1.0 / (qtype_traits_max * 512.0) + + qtype_max = float(qtype_traits_max) + + input = input.view(num_tokens, -1, group_size) + out = out.view(num_tokens, -1, group_size) + + for tile_m, tile_gn, tile_n in hl.tile( + [num_tokens, groups_per_row, group_size], block_size=[1, None, group_size] + ): + x_a_blk = input[tile_m, tile_gn, tile_n].to(torch.float32) + x_b_blk = hl.load( + input, + [tile_m, tile_gn.index + groups_per_row, tile_n], + extra_mask=(tile_gn.index + groups_per_row < 2 * groups_per_row)[ + None, :, None + ], + ).to(torch.float32) + x_blk = x_a_blk * torch.sigmoid(x_a_blk) * x_b_blk + s_blk = torch.amax(torch.abs(x_blk), dim=-1).to(torch.float32) + + if scale_ub is not None: + scale_ub_s = hl.load(scale_ub, []) + s_blk = s_blk.clamp(max=scale_ub_s) + s_blk = s_blk * (1.0 / qtype_max) + s_blk = s_blk.clamp(min=min_scaling_factor) + + scales[tile_m, tile_gn] = s_blk + if quant_dtype == torch.int8: + y_blk = (x_blk * (1.0 / s_blk[:, :, None])).round() + else: + y_blk = x_blk / s_blk[:, :, None] + + out[tile_m, tile_gn, tile_n] = y_blk.clamp( + qtype_traits_min, qtype_traits_max + ).to(out.dtype) diff --git a/vllm/lora/layers/base_linear.py b/vllm/lora/layers/base_linear.py index bff3c0cf454..5c8e829b299 100644 --- a/vllm/lora/layers/base_linear.py +++ b/vllm/lora/layers/base_linear.py @@ -89,7 +89,7 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA): vllm_config = get_current_vllm_config() self._lora_stream = _get_lora_aux_cuda_stream() assert current_platform.is_cuda_alike() - self._events = [torch.Event(), torch.Event()] + self._events = [torch.cuda.Event(), torch.cuda.Event()] # lora_linear avoids prefix conflicts with the base layer self.layer_name = self.base_layer.prefix + ".lora_linear_async" compilation_config = vllm_config.compilation_config diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index 39f60aad5db..63a4ea9a829 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -118,7 +118,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA): def _init_lora_stream_context(self) -> None: self._lora_stream: torch.cuda.Stream | None = None - self._events: tuple[torch.Event, ...] | None = None + self._events: tuple[torch.cuda.Event, ...] | None = None if not self._enable_aux_cuda_stream: return if not current_platform.is_cuda_alike(): @@ -127,7 +127,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA): # 4 events: 2 per (base GEMM, LoRA) pair so w13 and w2 don't reuse # the same event objects; reuse-within-a-pair is fine because the # second pair starts only after intermediate_cache1.add_() has joined. - self._events = tuple(torch.Event() for _ in range(4)) + self._events = tuple(torch.cuda.Event() for _ in range(4)) def _build_lora_context(self): use_dual_stream = ( diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index f5e3a16d71b..5d1d3707308 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -74,6 +74,9 @@ from vllm.model_executor.kernels.linear.mxfp4 import ( from vllm.model_executor.kernels.linear.mxfp4.flashinfer import ( FlashInferMxFp4LinearKernel, ) +from vllm.model_executor.kernels.linear.mxfp4.humming import ( + HummingMxFp4LinearKernel, +) from vllm.model_executor.kernels.linear.mxfp4.marlin import ( MarlinMxFp4LinearKernel, ) @@ -91,6 +94,9 @@ from vllm.model_executor.kernels.linear.mxfp8.flashinfer import ( FlashInferCutedslMxfp8LinearKernel, FlashInferCutlassMxfp8LinearKernel, ) +from vllm.model_executor.kernels.linear.mxfp8.humming import ( + HummingMxfp8LinearKernel, +) from vllm.model_executor.kernels.linear.mxfp8.marlin import ( MarlinMxfp8LinearKernel, ) @@ -120,6 +126,9 @@ from vllm.model_executor.kernels.linear.nvfp4.flashinfer import ( FlashInferCutlassNvFp4LinearKernel, FlashInferTrtllmNvFp4LinearKernel, ) +from vllm.model_executor.kernels.linear.nvfp4.humming import ( + HummingNvFp4LinearKernel, +) from vllm.model_executor.kernels.linear.nvfp4.marlin import ( MarlinNvFp4LinearKernel, ) @@ -154,6 +163,10 @@ from vllm.model_executor.kernels.linear.scaled_mm.flashinfer import ( FlashInferFp8DeepGEMMDynamicBlockScaledKernel, FlashInferFP8ScaledMMLinearKernel, ) +from vllm.model_executor.kernels.linear.scaled_mm.humming import ( + HummingFP8ScaledMMLinearKernel, + HummingInt8ScaledMMLinearKernel, +) from vllm.model_executor.kernels.linear.scaled_mm.marlin import ( MarlinFP8ScaledMMLinearKernel, ) @@ -225,6 +238,14 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = { "flashinfer_b12x": { FlashInferB12xNvFp4LinearKernel, }, + "humming": { + HummingFP8ScaledMMLinearKernel, + HummingInt8ScaledMMLinearKernel, + HummingLinearKernel, + HummingMxfp8LinearKernel, + HummingMxFp4LinearKernel, + HummingNvFp4LinearKernel, + }, "marlin": { MarlinFP8ScaledMMLinearKernel, MarlinLinearKernel, @@ -292,6 +313,7 @@ _POSSIBLE_INT8_KERNELS: dict[PlatformEnum, list[type[Int8ScaledMMLinearKernel]]] PlatformEnum.CUDA: [ CutlassInt8ScaledMMLinearKernel, TritonInt8ScaledMMLinearKernel, + HummingInt8ScaledMMLinearKernel, ], PlatformEnum.ROCM: [AiterInt8ScaledMMLinearKernel, TritonInt8ScaledMMLinearKernel], } @@ -304,6 +326,7 @@ _POSSIBLE_FP8_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] = CutlassFP8ScaledMMLinearKernel, PerTensorTorchFP8ScaledMMLinearKernel, ChannelWiseTorchFP8ScaledMMLinearKernel, + HummingFP8ScaledMMLinearKernel, ], PlatformEnum.ROCM: [ AiterHipbMMPerTokenFp8ScaledMMLinearKernel, @@ -335,6 +358,7 @@ _POSSIBLE_FP8_BLOCK_KERNELS: dict[ CutlassFp8BlockScaledMMKernel, MarlinFP8ScaledMMLinearKernel, TritonFp8BlockScaledMMKernel, + HummingFP8ScaledMMLinearKernel, ], PlatformEnum.ROCM: [ AiterFp8BlockScaledMMKernel, @@ -351,6 +375,7 @@ _POSSIBLE_FP8_BLOCK_KERNELS: dict[ _POSSIBLE_WFP8A16_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] = { PlatformEnum.CUDA: [ + HummingFP8ScaledMMLinearKernel, MarlinFP8ScaledMMLinearKernel, ], PlatformEnum.ROCM: [ @@ -371,10 +396,10 @@ _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = { MacheteLinearKernel, AllSparkLinearKernel, MarlinLinearKernel, - HummingLinearKernel, ConchLinearKernel, ExllamaLinearKernel, TritonW4A16LinearKernel, + HummingLinearKernel, ], PlatformEnum.ROCM: [ RDNA3W4A16LinearKernel, @@ -400,6 +425,7 @@ _POSSIBLE_MXFP8_KERNELS: dict[PlatformEnum, list[type[Mxfp8LinearKernel]]] = { FlashInferCutlassMxfp8LinearKernel, MarlinMxfp8LinearKernel, EmulationMxfp8LinearKernel, + HummingMxfp8LinearKernel, ], PlatformEnum.ROCM: [ # Native CDNA4 (gfx950) MX linear; is_supported() gates to gfx95x and @@ -426,6 +452,7 @@ _POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = { FlashInferCudnnNvFp4LinearKernel, FbgemmNvFp4LinearKernel, EmulationNvFp4LinearKernel, + HummingNvFp4LinearKernel, ], PlatformEnum.ROCM: [ EmulationNvFp4LinearKernel, @@ -436,6 +463,7 @@ _POSSIBLE_MXFP4_KERNELS: dict[PlatformEnum, list[type[MxFp4LinearKernel]]] = { PlatformEnum.CUDA: [ FlashInferMxFp4LinearKernel, MarlinMxFp4LinearKernel, + HummingMxFp4LinearKernel, ], PlatformEnum.XPU: [ XPUMxFp4LinearKernel, diff --git a/vllm/model_executor/kernels/linear/mixed_precision/__init__.py b/vllm/model_executor/kernels/linear/mixed_precision/__init__.py index c0b8c35bbd5..f9b4bd6e435 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/__init__.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/__init__.py @@ -19,6 +19,9 @@ from vllm.model_executor.kernels.linear.mixed_precision.dynamic_4bit import ( from vllm.model_executor.kernels.linear.mixed_precision.exllama import ( ExllamaLinearKernel, ) +from vllm.model_executor.kernels.linear.mixed_precision.humming import ( + HummingLinearKernel, +) from vllm.model_executor.kernels.linear.mixed_precision.machete import ( MacheteLinearKernel, ) @@ -52,6 +55,7 @@ __all__ = [ "CutlassW4A8LinearKernel", "Dynamic4bitLinearKernel", "ExllamaLinearKernel", + "HummingLinearKernel", "MacheteLinearKernel", "MarlinLinearKernel", "RDNA3W4A16LinearKernel", diff --git a/vllm/model_executor/kernels/linear/mixed_precision/humming.py b/vllm/model_executor/kernels/linear/mixed_precision/humming.py index 764c0f4227f..7f1a9024ee1 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/humming.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/humming.py @@ -23,8 +23,6 @@ class HummingLinearKernel(MPLinearKernel): return False, "Humming is not installed" if c.has_g_idx: return False, "Humming does not support act-order (g_idx)" - if c.zero_points: - return False, "Humming linear kernel only supports symmetric weights" return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: @@ -41,6 +39,11 @@ class HummingLinearKernel(MPLinearKernel): "group_size": 0 if group_size == -1 else group_size, } + if self.config.zero_points: + assert self.w_zp_name is not None + name_map["zero_point"] = self.w_zp_name + quant_config["has_zero_point"] = True + convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map) prepare_humming_layer(layer, quant_config) diff --git a/vllm/model_executor/kernels/linear/mxfp4/humming.py b/vllm/model_executor/kernels/linear/mxfp4/humming.py new file mode 100644 index 00000000000..d93f5d48158 --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp4/humming.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.model_executor.layers.quantization.utils.humming_utils import ( + convert_linear_layer_to_humming_standard, + prepare_humming_layer, +) +from vllm.platforms import current_platform + +from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig + + +class HummingMxFp4LinearKernel(MxFp4LinearKernel): + """Humming GEMM Kernel for MXFP4.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_cuda(): + return False, "Humming only supported on CUDA" + + if not current_platform.has_device_capability(75): + return False, "Humming only supported on SM75+" + + return True, None + + @classmethod + def can_implement(cls, c: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.weight_scale.data = layer.weight_scale.data.view(torch.float8_e8m0fnu) + name_map = {"weight": "weight", "weight_scale": "weight_scale"} + + quant_config = { + "quant_method": "humming", + "dtype": "float4e2m1", + "scale_dtype": "float8e8m0", + "group_size": 32, + "weight_scale_type": "group", + } + + convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map) + prepare_humming_layer(layer, quant_config) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + from vllm.utils.humming import HummingMethod + + flatten_inputs = x.view(-1, x.size(-1)) + output = HummingMethod.forward_layer( + layer=layer, + inputs=flatten_inputs, + compute_config=layer.compute_config, + ) + return output.view(*x.shape[:-1], output.size(-1)) diff --git a/vllm/model_executor/kernels/linear/mxfp8/humming.py b/vllm/model_executor/kernels/linear/mxfp8/humming.py new file mode 100644 index 00000000000..ed1cd39cbd6 --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp8/humming.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.model_executor.layers.quantization.utils.humming_utils import ( + convert_linear_layer_to_humming_standard, + prepare_humming_layer, +) +from vllm.platforms import current_platform + +from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig + + +class HummingMxfp8LinearKernel(Mxfp8LinearKernel): + """Humming GEMM Kernel for MXFP8.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_cuda(): + return False, "Humming only supported on CUDA" + + if not current_platform.has_device_capability(75): + return False, "Humming only supported on SM75+" + + return True, None + + @classmethod + def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.weight_scale.data = layer.weight_scale.data.view(torch.float8_e8m0fnu) + name_map = {"weight": "weight", "weight_scale": "weight_scale"} + + quant_config = { + "quant_method": "humming", + "dtype": "float8e4m3", + "scale_dtype": "float8e8m0", + "group_size": 32, + "weight_scale_type": "group", + } + + convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map) + prepare_humming_layer(layer, quant_config) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + from vllm.utils.humming import HummingMethod + + flatten_inputs = x.view(-1, x.size(-1)) + output = HummingMethod.forward_layer( + layer=layer, + inputs=flatten_inputs, + compute_config=layer.compute_config, + ) + return output.view(*x.shape[:-1], output.size(-1)) diff --git a/vllm/model_executor/kernels/linear/nvfp4/humming.py b/vllm/model_executor/kernels/linear/nvfp4/humming.py new file mode 100644 index 00000000000..7e390343854 --- /dev/null +++ b/vllm/model_executor/kernels/linear/nvfp4/humming.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.utils.humming_utils import ( + prepare_humming_layer, +) +from vllm.platforms import current_platform + +from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig + +logger = init_logger(__name__) + + +class HummingNvFp4LinearKernel(NvFp4LinearKernel): + """Humming GEMM Kernel for NVFP4.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_cuda(): + return False, "Humming only supported on CUDA" + + if not current_platform.has_device_capability(75): + return False, "Humming only supported on SM75+" + + return True, None + + @classmethod + def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # Route through humming's compressed-tensors nvfp4 loader (same path as + # the MoE oracle); the native group_tensor schema mishandles a scalar + # global scale. + quant_config = { + "quant_method": "compressed-tensors", + "format": "nvfp4-pack-quantized", + "type": "float", + "num_bits": 4, + "strategy": "group", + "group_size": 16, + } + # CT pack-quantized reads `weight_packed`; the scheme renamed it to `weight`. + if not hasattr(layer, "weight_packed"): + layer.weight_packed = layer.weight + del layer.weight + # The CT linear scheme inverts the global scale (1/scale) for + # marlin/cutlass; humming wants the original. + layer.weight_global_scale = torch.nn.Parameter( + 1.0 / layer.weight_global_scale, requires_grad=False + ) + prepare_humming_layer(layer, quant_config) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + from vllm.utils.humming import HummingMethod + + flatten_inputs = x.view(-1, x.size(-1)) + output = HummingMethod.forward_layer( + layer=layer, + inputs=flatten_inputs, + compute_config=layer.compute_config, + ) + return output.view(*x.shape[:-1], output.size(-1)) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/humming.py b/vllm/model_executor/kernels/linear/scaled_mm/humming.py new file mode 100644 index 00000000000..7b8ed21fbd6 --- /dev/null +++ b/vllm/model_executor/kernels/linear/scaled_mm/humming.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.utils.humming_utils import ( + convert_linear_layer_to_humming_standard, + prepare_humming_layer, +) +from vllm.platforms import current_platform + +from .ScaledMMLinearKernel import ( + FP8ScaledMMLinearKernel, + FP8ScaledMMLinearLayerConfig, + Int8ScaledMMLinearKernel, + Int8ScaledMMLinearLayerConfig, +) + +logger = init_logger(__name__) + + +class HummingFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): + """Humming GEMM Kernel for FP8.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_cuda(): + return False, "Humming only supported on CUDA" + + if not current_platform.has_device_capability(75): + return False, "Humming only supported on SM75+" + + return True, None + + @classmethod + def can_implement( + cls, config: FP8ScaledMMLinearLayerConfig + ) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + from vllm.utils.humming import dtypes + + name_map = {"weight": "weight", "weight_scale": "weight_scale"} + scale_torch_dtype = self.config.weight_quant_key.scale.dtype + scale_dtype = dtypes.DataType.from_torch_dtype(scale_torch_dtype) + + quant_config = { + "quant_method": "humming", + "dtype": "float8e4m3", + "scale_dtype": scale_dtype, + } + + assert self.config.weight_quant_key.scale2 is None + scale_group_shape = self.config.weight_quant_key.scale.group_shape + if scale_group_shape.is_per_tensor(): + quant_config["weight_scale_type"] = "tensor" + if not hasattr(layer, "global_scale") and hasattr(layer, "weight_scale"): + del name_map["weight_scale"] + name_map["global_scale"] = "weight_scale" + elif scale_group_shape.is_per_channel(): + quant_config["weight_scale_type"] = "channel" + elif scale_group_shape.is_per_group(): + quant_config["weight_scale_type"] = "group" + quant_config["group_size"] = scale_group_shape.col + else: + assert scale_group_shape.row > 0 and scale_group_shape.col > 0 + quant_config["weight_scale_type"] = "block" + quant_config["weight_scale_group_size_n"] = scale_group_shape.row + quant_config["weight_scale_group_size"] = scale_group_shape.col + + if hasattr(layer, "weight_scale_inv"): + name_map["weight_scale"] = "weight_scale_inv" + + convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map) + prepare_humming_layer(layer, quant_config) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + from vllm.utils.humming import HummingMethod + + flatten_inputs = x.view(-1, x.size(-1)) + output = HummingMethod.forward_layer( + layer=layer, + inputs=flatten_inputs, + compute_config=layer.compute_config, + ) + return output.view(*x.shape[:-1], output.size(-1)) + + def apply_scaled_mm( + self, + *, + A: torch.Tensor, + B: torch.Tensor, + out_dtype: torch.dtype, + As: torch.Tensor, + Bs: torch.Tensor, + bias: torch.Tensor | None, + output_shape: list, + ) -> torch.Tensor: + pass + + +class HummingInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel): + """Humming GEMM Kernel for INT8.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_cuda(): + return False, "Humming only supported on CUDA" + + if not current_platform.has_device_capability(75): + return False, "Humming only supported on SM75+" + + return True, None + + @classmethod + def can_implement( + cls, config: Int8ScaledMMLinearLayerConfig + ) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + weight_name, weight_scale_name, *_ = self.layer_param_names + name_map = {"weight": weight_name, "weight_scale": weight_scale_name} + quant_config = {"quant_method": "humming", "dtype": "int8"} + weight = getattr(layer, weight_name) + weight.data = weight.data + 128 + + convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map) + prepare_humming_layer(layer, quant_config) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + from vllm.utils.humming import HummingMethod + + flatten_inputs = x.view(-1, x.size(-1)) + output = HummingMethod.forward_layer( + layer=layer, + inputs=flatten_inputs, + compute_config=layer.compute_config, + ) + return output.view(*x.shape[:-1], output.size(-1)) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/xpu.py b/vllm/model_executor/kernels/linear/scaled_mm/xpu.py index 3d23b02bea8..f30d6ced1d3 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/xpu.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/xpu.py @@ -197,29 +197,6 @@ class XPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): return False, "XPUFp8BlockScaledMM only support on XPU" return True, None - def process_weights_after_loading(self, layer: torch.nn.Module): - super().process_weights_after_loading(layer) - scale_attr = ( - "weight_scale_inv" if hasattr(layer, "weight_scale_inv") else "weight_scale" - ) - scale = getattr(layer, scale_attr) - # Models with scale_fmt=ue8m0 (e.g. DeepSeek-V4) store weight scales - # as float8_e8m0fnu. The oneDNN fp8_gemm kernel dispatches to its - # "block quant" path only when NEITHER scale is e8m0: - # - # is_block_quant = (m1_sc != e8m0) && (m2_sc != e8m0) && ... - # - # Since activation scales are always float32 (use_ue8m0=False on XPU, - # DeepGEMM requires Hopper/Blackwell), an e8m0 weight scale causes - # is_block_quant=false and falls into the wrong per-channel path, - # producing NaN. Converting e8m0→float32 here at load time (one-time, - # negligible overhead for small scale tensors) ensures the kernel sees - # matching dtypes and correctly enters the block-quant path with the - # actual group_size derived from scale tensor shapes. - if scale.dtype == torch.float8_e8m0fnu: - scale = scale.to(torch.float32) - replace_parameter(layer, scale_attr, scale.data.t().contiguous()) - def apply_block_scaled_mm( self, A: torch.Tensor, @@ -228,11 +205,12 @@ class XPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): Bs: torch.Tensor, ) -> torch.Tensor: # Weight is [N, K]. Use .t() to create a [K, N] view without copying. + # Bs is [N/128, K/128] — transpose to [K/128, N/128] for oneDNN. return torch.ops._xpu_C.fp8_gemm( A, B.t(), self.config.out_dtype, As, - Bs, + Bs.t().contiguous(), torch.Tensor(), ) diff --git a/vllm/model_executor/layers/fla/ops/layernorm_guard.py b/vllm/model_executor/layers/fla/ops/layernorm_guard.py index 8b9e275737e..279b06b5ebc 100644 --- a/vllm/model_executor/layers/fla/ops/layernorm_guard.py +++ b/vllm/model_executor/layers/fla/ops/layernorm_guard.py @@ -331,46 +331,6 @@ def rmsnorm_fn( ) -class LayerNormGated(nn.Module): - def __init__( - self, - hidden_size, - eps: float = 1e-5, - group_size: int | None = None, - norm_before_gate: bool = True, - device: torch.device | None = None, - dtype: torch.dtype | None = None, - ): - """If group_size is not None, we do GroupNorm with each group having group_size elements. - group_size=None is equivalent to group_size=hidden_size (i.e. there's only 1 group). - """ - - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.bias = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.group_size = group_size - self.norm_before_gate = norm_before_gate - self.reset_parameters() - - def reset_parameters(self): - torch.nn.init.ones_(self.weight) - torch.nn.init.zeros_(self.bias) - - def forward(self, x, z=None): - """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z))""" - return layernorm_fn( - x, - self.weight, - self.bias, - z=z, - group_size=self.group_size, - eps=self.eps, - norm_before_gate=self.norm_before_gate, - ) - - class RMSNormGated(nn.Module): def __init__( self, 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 c21072676ec..f4933e1bd76 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 @@ -14,12 +14,13 @@ from vllm.model_executor.layers.fused_moe.config import ( ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, + kInt4W4A8StaticChannelSym, kInt4W4A8StaticGroup32Sym, kInt4W4A8StaticGroup64Sym, kInt4W4A8StaticGroup128Sym, kInt4W4A8StaticGroupSym, ) -from vllm.platforms import current_platform +from vllm.platforms import CpuArchEnum, current_platform class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic): @@ -48,6 +49,48 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic): def activation_format() -> mk.FusedMoEActivationFormat: return mk.FusedMoEActivationFormat.Standard + @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]: + if ( + not current_platform.is_cpu() + or current_platform.get_cpu_architecture() != CpuArchEnum.ARM + ): + return False, "kernel only supports Arm CPU" + + if moe_config.in_dtype not in ( + torch.float32, + torch.bfloat16, + torch.float16, + ): + return ( + False, + f"kernel does not support {moe_config.in_dtype} input/output dtype", + ) + + try: + _ = torch.ops.aten._dyn_quant_matmul_4bit + _ = torch.ops.aten._dyn_quant_pack_4bit_weight + except AttributeError: + return ( + False, + f"PyTorch {torch.__version__} does not support " + "_dyn_quant_* 4bit ops. Install a newer version", + ) + + return mk.FusedMoEExperts.is_supported_config( + cls, + moe_config, + weight_key, + activation_key, + activation_format, + ) + @staticmethod def _supports_current_device() -> bool: return current_platform.is_cpu() @@ -86,8 +129,9 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic): Can be channel-wise or group-wise quantization - Activations: dynamic per-token 8-bit integer quantization """ - # group size must be multiple of 32 + # channelwise or groupwise with group size being a multiple of 32 SUPPORTED_W_A = [ + (kInt4W4A8StaticChannelSym, None), (kInt4W4A8StaticGroup128Sym, None), (kInt4W4A8StaticGroup64Sym, None), (kInt4W4A8StaticGroup32Sym, None), @@ -120,7 +164,8 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic): """Expert parallelism not yet supported.""" return False - def _activation_kind(self, activation: MoEActivation) -> int: + @staticmethod + def _activation_kind(activation: MoEActivation) -> int: """Convert MoEActivation to kernel activation kind integer. Returns: @@ -200,29 +245,22 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic): e_score_correction_bias=e_score_correction_bias, ) - # Extract dimensions from weight tensors - # w1 is w13_packed: [num_experts, packed_data...] - # w2 is w2_packed: [num_experts, packed_data...] - # These dimensions should be available from the layer - # For now, we'll extract from moe_config - K = self.moe_config.hidden_dim - N = self.moe_config.intermediate_size_per_partition + hidden_size = self.moe_config.hidden_dim + intermediate_size = self.moe_config.intermediate_size_per_partition assert self.quant_config.block_shape is not None - if self.quant_config.is_per_act_token: + # C++ kernel expects an int: -1 for channelwise, and group size for groupwise + if self.quant_config.block_shape == [-1, 1]: group_size = -1 else: group_size = self.quant_config.block_shape[1] - - # Call the dynamic 4-bit int MoE kernel return torch.ops._C.dynamic_4bit_int_moe( hidden_states, topk_ids.to(torch.long), topk_weights, w1, # w13_weight_packed w2, # w2_weight_packed - K, # hidden_size (w2_out_features) - N, # intermediate_size (w2_in_features) - N * 2, # 2*intermediate_size (w13_out_features) + hidden_size, + intermediate_size, group_size, apply_router_weight_on_input, self._activation_kind(activation), diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_b12x_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_b12x_moe.py index 38200d9d090..ca23e127249 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_b12x_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_b12x_moe.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any + import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk @@ -20,7 +22,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( ) from vllm.platforms import current_platform from vllm.utils.flashinfer import ( - flashinfer_b12x_fused_moe, flashinfer_convert_sf_to_mma_layout, has_flashinfer_b12x_moe, ) @@ -42,6 +43,11 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular): Only NVFP4 (kNvfp4Static/kNvfp4Dynamic) quantization is supported. """ + _ACTIVATION_MAP: dict[MoEActivation, str] = { + MoEActivation.SILU: "silu", + MoEActivation.RELU2_NO_MUL: "relu2", + } + def __init__( self, moe_config: FusedMoEConfig, @@ -60,6 +66,30 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular): # one. Holding it on the instance keeps apply() alloc-free. self._fc2_input_scale: torch.Tensor | None = None + # Shape params for B12xMoEWrapper construction. + self.global_num_experts = moe_config.num_experts + self.topk = moe_config.experts_per_token + self.hidden_dim = moe_config.hidden_dim + self.intermediate_size_per_partition = ( + moe_config.intermediate_size_per_partition + ) + self.max_num_tokens = moe_config.max_num_tokens + self.local_expert_offset = self.ep_rank * self.num_local_experts + + activation = moe_config.activation + if activation not in self._ACTIVATION_MAP: + raise ValueError( + f"FlashInferB12xExperts does not support " + f"activation {activation!r}. " + f"Supported: {list(self._ACTIVATION_MAP.keys())}" + ) + self._activation_str = self._ACTIVATION_MAP[activation] + + # Lazily created on first apply() call. + self._wrapper: Any | None = None + self.w1_sf_mma: torch.Tensor | None = None + self.w2_sf_mma: torch.Tensor | None = None + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # Normalise block scales to absorb the per-expert weight global scale # (w_gs). vLLM's NVFP4 convention stores: @@ -141,7 +171,7 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular): @staticmethod def _supports_no_act_and_mul() -> bool: - return False + return True @staticmethod def _supports_quant_scheme( @@ -158,11 +188,13 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular): @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - return activation == MoEActivation.SILU + return activation in (MoEActivation.SILU, MoEActivation.RELU2_NO_MUL) @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: - return True + # B12xMoEWrapper does not yet support expert parallelism: its local + # expert count must equal the global expert count. + return not moe_parallel_config.use_ep def supports_expert_map(self) -> bool: return False @@ -190,13 +222,29 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular): @property def expects_unquantized_inputs(self) -> bool: - # b12x_fused_moe expects BF16 hidden states and performs its own FP4 + # B12xMoEWrapper expects BF16 hidden states and performs its own FP4 # quantization internally. Returning True prevents the modular kernel - # from pre-quantizing activations, which would produce an FP4-packed - # tensor with size(-1)=k//2 and break the scale-factor conversion that - # expects size(-1)=k. + # from pre-quantizing activations. return True + def _ensure_wrapper(self) -> None: + """Lazily create B12xMoEWrapper on first use.""" + if self._wrapper is not None: + return + + from flashinfer.fused_moe import B12xMoEWrapper + + self._wrapper = B12xMoEWrapper( + num_experts=self.global_num_experts, + top_k=self.topk, + hidden_size=self.hidden_dim, + intermediate_size=self.intermediate_size_per_partition, + use_cuda_graph=True, + max_num_tokens=self.max_num_tokens, + num_local_experts=self.num_local_experts, + activation=self._activation_str, + ) + def apply( self, output: torch.Tensor, @@ -224,13 +272,16 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular): assert self._fc2_input_scale is not None, ( "_fc2_input_scale must be set by process_weights_after_loading" ) + assert self.w1_sf_mma is not None and self.w2_sf_mma is not None, ( + "process_weights_after_loading must run before FlashInferB12xExperts.apply" + ) - top_k = topk_ids.shape[1] + self._ensure_wrapper() + wrapper = self._wrapper + assert wrapper is not None - flashinfer_b12x_fused_moe( + wrapper_output = wrapper.run( x=hidden_states, - token_selected_experts=topk_ids.to(torch.int32), - token_final_scales=topk_weights, w1_weight=w1, w1_weight_sf=self.w1_sf_mma, w1_alpha=self.g1_alphas, @@ -238,9 +289,7 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular): w2_weight=w2, w2_weight_sf=self.w2_sf_mma, w2_alpha=self.g2_alphas, - num_experts=global_num_experts, - top_k=top_k, - num_local_experts=self.num_local_experts, - output_dtype=self.out_dtype, - output=output, + token_selected_experts=topk_ids.to(torch.int32), + token_final_scales=topk_weights, ) + output.copy_(wrapper_output) diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py index 047ae46c0d3..5f112380cf9 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py @@ -38,15 +38,20 @@ from vllm.model_executor.layers.fused_moe.utils import ( ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, + kFp8Dynamic128Sym, kFp8DynamicTokenSym, kFp8Static128BlockSym, kFp8StaticChannelSym, + kFp8StaticTensorSym, kInt4Static, + kInt8DynamicTokenSym, kInt8Static, + kInt8StaticChannelSym, kMxfp4Dynamic, kMxfp4Static, kMxfp8Dynamic, kMxfp8Static, + kNvfp4Dynamic, kNvfp4Static, ) from vllm.platforms import current_platform @@ -61,9 +66,9 @@ if TYPE_CHECKING: logger = init_logger(__name__) -def get_humming_moe_gemm_type() -> str | None: +def get_humming_moe_gemm_type() -> str: env_gemm_type: str | None = envs.VLLM_HUMMING_MOE_GEMM_TYPE - gemm_type = None + gemm_type = "indexed" if env_gemm_type is not None: env_gemm_type = env_gemm_type.lower() if env_gemm_type == "indexed": @@ -87,7 +92,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): num_dispatchers: int | None = None, ): self.layer = layer - self.num_experts = self.layer.num_experts + self.num_experts = self.layer.local_num_experts self.global_num_experts = self.layer.global_num_experts self.init_humming_moe() @@ -186,6 +191,24 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): (kInt4Static, kFp8DynamicTokenSym), (kInt8Static, None), (kInt8Static, kFp8DynamicTokenSym), + # Checkpoint-driven (weight, activation) pairs the dense/MoE oracles + # pass. Humming defers input quant (see expects_unquantized_inputs), + # so the activation key does not constrain support. + # fp8 (compressed-tensors / native / modelopt) + (kFp8StaticChannelSym, kFp8StaticTensorSym), + (kFp8StaticChannelSym, kFp8Dynamic128Sym), + (kFp8StaticTensorSym, None), + (kFp8StaticTensorSym, kFp8DynamicTokenSym), + (kFp8StaticTensorSym, kFp8StaticTensorSym), + (kFp8StaticTensorSym, kFp8Dynamic128Sym), + (kFp8Static128BlockSym, kFp8Dynamic128Sym), + # int8 (compressed-tensors w8a8 / experts_int8) + (kInt8StaticChannelSym, None), + (kInt8StaticChannelSym, kInt8DynamicTokenSym), + # nvfp4 (compressed-tensors / modelopt / quark) + (kNvfp4Static, kNvfp4Dynamic), + # mxfp8 (compressed-tensors / modelopt / online) + (kMxfp8Static, kMxfp8Dynamic), ] return (weight_key, activation_key) in SUPPORTED_W_A @@ -463,13 +486,12 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): assert hasattr(cls, "humming_gemm_type") gemm_type = cls.humming_gemm_type().value.lower() preferred_gemm_type = get_humming_moe_gemm_type() - if preferred_gemm_type is not None: - supported = preferred_gemm_type.lower() == gemm_type - if not supported: - reason = ( - f"preferred gemm type {preferred_gemm_type} != " - f"supported gemm type {gemm_type}" - ) + supported = preferred_gemm_type.lower() == gemm_type + if not supported: + reason = ( + f"preferred gemm type {preferred_gemm_type} != " + f"supported gemm type {gemm_type}" + ) return supported, reason diff --git a/vllm/model_executor/layers/fused_moe/experts/lora_context.py b/vllm/model_executor/layers/fused_moe/experts/lora_context.py index dd7429d76e1..117f744aeea 100644 --- a/vllm/model_executor/layers/fused_moe/experts/lora_context.py +++ b/vllm/model_executor/layers/fused_moe/experts/lora_context.py @@ -51,7 +51,7 @@ class MoELoRAContext: # Events are paired one-per-overlap-pair: events[0,1] for w13, # events[2,3] for w2, so the two pairs do not race on the same event. aux_stream: torch.cuda.Stream | None = None - events: tuple[torch.Event, ...] | None = None + events: tuple[torch.cuda.Event, ...] | None = None # Per-rank token→LoRA mapping after EP dispatch. Set by # FusedMoEPrepareAndFinalizeModular.prepare() when EP+LoRA is active, read diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 862f292009c..c60e68232d4 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from enum import Enum +from typing import Any import torch @@ -45,6 +46,7 @@ class Fp8MoeBackend(Enum): DEEPGEMM = "DEEPGEMM" BATCHED_DEEPGEMM = "BATCHED_DEEPGEMM" MARLIN = "MARLIN" + HUMMING = "HUMMING" TRITON = "TRITON" BATCHED_TRITON = "BATCHED_TRITON" AITER = "AITER" @@ -83,6 +85,7 @@ def _get_priority_backends( Fp8MoeBackend.VLLM_CUTLASS, Fp8MoeBackend.TRITON, Fp8MoeBackend.MARLIN, + Fp8MoeBackend.HUMMING, Fp8MoeBackend.BATCHED_DEEPGEMM, Fp8MoeBackend.BATCHED_VLLM_CUTLASS, Fp8MoeBackend.BATCHED_TRITON, @@ -162,6 +165,19 @@ def backend_to_kernel_cls( return [BatchedDeepGemmExperts] + elif backend == Fp8MoeBackend.HUMMING: + from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( + BatchedHummingGroupedExperts, + HummingGroupedExperts, + HummingIndexedExperts, + ) + + return [ + BatchedHummingGroupedExperts, + HummingGroupedExperts, + HummingIndexedExperts, + ] + elif backend == Fp8MoeBackend.MARLIN: from vllm.model_executor.layers.fused_moe.experts.marlin_moe import ( MarlinExperts, @@ -240,6 +256,7 @@ def map_fp8_backend(runner_backend: MoEBackend) -> Fp8MoeBackend: "flashinfer_trtllm": Fp8MoeBackend.FLASHINFER_TRTLLM, "flashinfer_cutlass": Fp8MoeBackend.FLASHINFER_CUTLASS, "marlin": Fp8MoeBackend.MARLIN, + "humming": Fp8MoeBackend.HUMMING, "aiter": Fp8MoeBackend.AITER, "hpc": Fp8MoeBackend.HPC, } @@ -402,6 +419,41 @@ def select_fp8_moe_backend( return Fp8MoeBackend.NONE, None +def _humming_fp8_weight_schema( + layer: RoutedExperts, weight: torch.Tensor, weight_scale: torch.Tensor +) -> dict[str, Any]: + """Build the humming weight schema from the canonical on-device fp8/mxfp8 + tensors (scale dtype/shape, block size), not the producing quant method.""" + # mxfp8: e8m0 group-32 scales (stored as uint8 bytes or e8m0). humming has + # no compressed-tensors mxfp8 loader; its modelopt schema fits both sources. + if weight_scale.dtype in (torch.uint8, torch.float8_e8m0fnu): + return {"quant_method": "modelopt", "quant_algo": "mxfp8"} + + if hasattr(layer, "w13_weight_scale_inv"): + assert hasattr(layer, "weight_block_size") + return {"quant_method": "fp8", "weight_block_size": layer.weight_block_size} + + # fp8 (e4m3): recover the strategy from the scale layout (block from + # weight_block_size, else channel vs tensor by per-expert scale count). + config: dict[str, Any] = { + "quant_method": "compressed-tensors", + "format": "float-quantized", + "type": "float", + "num_bits": 8, + "symmetric": True, + } + weight_block_size = getattr(layer, "weight_block_size", None) + num_experts, num_output = weight.shape[0], weight.shape[-2] + if weight_block_size is not None: + config["strategy"] = "block" + config["block_structure"] = list(weight_block_size) + elif weight_scale.numel() >= num_experts * num_output: + config["strategy"] = "channel" + else: + config["strategy"] = "tensor" + return config + + def convert_to_fp8_moe_kernel_format( fp8_backend: Fp8MoeBackend, # TODO(bnell): replace layer with weight_block_size @@ -429,6 +481,18 @@ def convert_to_fp8_moe_kernel_format( w13, w2, w13_scale, w2_scale = rocm_aiter_ops.shuffle_mxfp8_moe_weights( w13, w2, w13_scale, w2_scale ) + elif fp8_backend == Fp8MoeBackend.HUMMING: + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + convert_to_humming_moe_kernel_format, + ) + + convert_to_humming_moe_kernel_format( + layer, quant_config=_humming_fp8_weight_schema(layer, w13, w13_scale) + ) + w13 = layer.w13_weight + w2 = layer.w2_weight + w13_scale = layer.w13_weight_scale + w2_scale = layer.w2_weight_scale elif fp8_backend == Fp8MoeBackend.MARLIN: weight_block_size = getattr(layer, "weight_block_size", None) if weight_block_size == [1, 32]: @@ -511,6 +575,7 @@ def make_fp8_moe_quant_config( swiglu_limit: float | None = None, gemm1_alpha: float | None = None, gemm1_beta: float | None = None, + layer: torch.nn.Module | None = None, ) -> FusedMoEQuantConfig: """ Create FusedMoEQuantConfig for the specified FP8 Backend. @@ -537,6 +602,14 @@ def make_fp8_moe_quant_config( gemm1_beta=gemm1_beta, gemm1_clamp_limit=swiglu_limit, ) + elif fp8_backend == Fp8MoeBackend.HUMMING: + from vllm.model_executor.layers.fused_moe import RoutedExperts + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + get_humming_moe_quant_config, + ) + + assert isinstance(layer, RoutedExperts) + return get_humming_moe_quant_config(layer) # Flashinfer CUTLASS or HPC per-tensor uses single dq scale # (alpha = w_scale * a_scale) and inverse a2 scale. @@ -600,6 +673,7 @@ def make_fp8_moe_kernel( experts_cls: type[mk.FusedMoEExperts], fp8_backend: Fp8MoeBackend, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, + layer: torch.nn.Module | None = None, ) -> mk.FusedMoEKernel: # Create Prepare/Finalize. prepare_finalize = maybe_make_prepare_finalize( @@ -613,6 +687,11 @@ def make_fp8_moe_kernel( logger.info_once("Using %s", prepare_finalize.__class__.__name__) + extra_kwargs = {} + if fp8_backend == Fp8MoeBackend.HUMMING: + assert layer is not None + extra_kwargs = {"layer": layer} + # Create Experts. if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts: max_num_tokens = prepare_finalize.max_num_tokens_per_rank() @@ -622,11 +701,13 @@ def make_fp8_moe_kernel( quant_config=moe_quant_config, max_num_tokens=max_num_tokens, num_dispatchers=prepare_finalize.num_dispatchers(), + **extra_kwargs, ) else: experts = experts_cls( moe_config=moe_config, quant_config=moe_quant_config, + **extra_kwargs, ) kernel = mk.FusedMoEKernel( diff --git a/vllm/model_executor/layers/fused_moe/oracle/int8.py b/vllm/model_executor/layers/fused_moe/oracle/int8.py index e31a3ca07ee..5a2b4c3a75b 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int8.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from enum import Enum +from typing import Any import torch @@ -22,6 +23,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kInt8DynamicTokenSym, kInt8StaticChannelSym, ) +from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform logger = init_logger(__name__) @@ -29,6 +31,7 @@ logger = init_logger(__name__) class Int8MoeBackend(Enum): TRITON = "TRITON" + HUMMING = "HUMMING" CPU = "CPU" @@ -40,6 +43,7 @@ def _get_priority_backends( """ _AVAILABLE_BACKENDS = [ Int8MoeBackend.TRITON, + Int8MoeBackend.HUMMING, Int8MoeBackend.CPU, ] @@ -62,6 +66,19 @@ def backend_to_kernel_cls( return [TritonExperts] + elif backend == Int8MoeBackend.HUMMING: + from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( + BatchedHummingGroupedExperts, + HummingGroupedExperts, + HummingIndexedExperts, + ) + + return [ + BatchedHummingGroupedExperts, + HummingGroupedExperts, + HummingIndexedExperts, + ] + elif backend == Int8MoeBackend.CPU: from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( CPUExpertsInt8, @@ -77,6 +94,7 @@ def map_int8_backend(runner_backend: MoEBackend) -> Int8MoeBackend: """Map user's MoEBackend to Int8MoeBackend.""" mapping = { "triton": Int8MoeBackend.TRITON, + "humming": Int8MoeBackend.HUMMING, } if backend := mapping.get(runner_backend): return backend @@ -163,6 +181,7 @@ def select_int8_moe_backend( def make_int8_moe_quant_config( + int8_backend: Int8MoeBackend, w1_scale: torch.Tensor, w2_scale: torch.Tensor, a1_scale: torch.Tensor | None = None, @@ -170,11 +189,21 @@ def make_int8_moe_quant_config( w1_bias: torch.Tensor | None = None, w2_bias: torch.Tensor | None = None, per_act_token_quant: bool = False, + layer: torch.nn.Module | None = None, ) -> FusedMoEQuantConfig: assert (a1_scale is None and a2_scale is None) or ( a1_scale is not None and a2_scale is not None ), "a1_scale and a2_scale must both be provided or both be None" + if int8_backend == Int8MoeBackend.HUMMING: + from vllm.model_executor.layers.fused_moe import RoutedExperts + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + get_humming_moe_quant_config, + ) + + assert isinstance(layer, RoutedExperts) + return get_humming_moe_quant_config(layer) + if a1_scale is None or a2_scale is None: return int8_w8a16_moe_quant_config( w1_scale=w1_scale, @@ -196,13 +225,56 @@ def make_int8_moe_quant_config( ) +def _humming_int8_weight_schema( + weight: torch.Tensor, weight_scale: torch.Tensor +) -> dict[str, Any]: + """Build the humming compressed-tensors int8 schema from the canonical + on-device tensors; humming does the signed-int8 -> native conversion.""" + config: dict[str, Any] = { + "quant_method": "compressed-tensors", + "format": "int-quantized", + "type": "int", + "num_bits": 8, + "symmetric": True, + "strategy": "channel", + } + num_experts, num_output = weight.shape[0], weight.shape[-2] + if weight_scale.numel() < num_experts * num_output: + config["strategy"] = "tensor" + return config + + def convert_to_int8_moe_kernel_format( int8_backend: Int8MoeBackend, w13: torch.Tensor, w2: torch.Tensor, + layer: torch.nn.Module | None = None, + w13_scale: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Convert INT8 MoE weights to backend-specific kernel format.""" - if int8_backend == Int8MoeBackend.CPU: + if int8_backend == Int8MoeBackend.HUMMING: + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + convert_to_humming_moe_kernel_format, + ) + + assert layer is not None + # Humming reads canonical CT scales (w*_weight_scale) from the layer. + # Online int8 produces per-channel (E, N) w*_scale; expose them as the + # (E, N, 1) w*_weight_scale humming's loader expects. + for sub in ("w13", "w2"): + if hasattr(layer, f"{sub}_weight_scale"): + continue + scale = getattr(layer, f"{sub}_scale").data + if scale.dim() < 3: + scale = scale.unsqueeze(-1) + replace_parameter(layer, f"{sub}_weight_scale", scale) + delattr(layer, f"{sub}_scale") + convert_to_humming_moe_kernel_format( + layer, + quant_config=_humming_int8_weight_schema(w13, layer.w13_weight_scale), + ) + return layer.w13_weight, layer.w2_weight + elif int8_backend == Int8MoeBackend.CPU: from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( prepare_int8_moe_layer_for_cpu, ) @@ -215,10 +287,12 @@ def convert_to_int8_moe_kernel_format( def make_int8_moe_kernel( + int8_backend: Int8MoeBackend, moe_quant_config: FusedMoEQuantConfig, moe_config: FusedMoEConfig, experts_cls: type[mk.FusedMoEExperts], routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, + layer: torch.nn.Module | None = None, ) -> mk.FusedMoEKernel: # Create Prepare/Finalize. prepare_finalize = maybe_make_prepare_finalize( @@ -232,6 +306,11 @@ def make_int8_moe_kernel( logger.info_once("Using %s", prepare_finalize.__class__.__name__) + extra_kwargs = {} + if int8_backend == Int8MoeBackend.HUMMING: + assert layer is not None + extra_kwargs = {"layer": layer} + # Create Experts. if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts: max_num_tokens = prepare_finalize.max_num_tokens_per_rank() @@ -241,11 +320,13 @@ def make_int8_moe_kernel( quant_config=moe_quant_config, max_num_tokens=max_num_tokens, num_dispatchers=prepare_finalize.num_dispatchers(), + **extra_kwargs, ) else: experts = experts_cls( moe_config=moe_config, quant_config=moe_quant_config, + **extra_kwargs, ) kernel = mk.FusedMoEKernel( diff --git a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py index 0cf1382d406..d0cc08ea141 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -11,6 +11,7 @@ from compressed_tensors.quantization import ( import vllm._custom_ops as ops import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.config.kernel import MoEBackend from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, @@ -46,6 +47,7 @@ logger = init_logger(__name__) class WNA16MoEBackend(Enum): MARLIN = "MARLIN" BATCHED_MARLIN = "BATCHED_MARLIN" + HUMMING = "HUMMING" CPU = "CPU" FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM" XPU = "XPU" @@ -55,7 +57,19 @@ def backend_to_kernel_cls( backend: WNA16MoEBackend, ) -> list[type[mk.FusedMoEExperts]]: """Return the experts class for the given backend, or None for NONE.""" - if backend == WNA16MoEBackend.MARLIN: + if backend == WNA16MoEBackend.HUMMING: + from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( + BatchedHummingGroupedExperts, + HummingGroupedExperts, + HummingIndexedExperts, + ) + + return [ + BatchedHummingGroupedExperts, + HummingGroupedExperts, + HummingIndexedExperts, + ] + elif backend == WNA16MoEBackend.MARLIN: return [MarlinExperts] elif backend == WNA16MoEBackend.BATCHED_MARLIN: return [BatchedMarlinExperts] @@ -90,10 +104,26 @@ def _get_priority_backends() -> list[WNA16MoEBackend]: WNA16MoEBackend.FLASHINFER_TRTLLM, WNA16MoEBackend.MARLIN, WNA16MoEBackend.BATCHED_MARLIN, + WNA16MoEBackend.HUMMING, ] return _AVAILABLE_BACKENDS +def map_wna16_backend(runner_backend: MoEBackend) -> WNA16MoEBackend: + """Map user's MoEBackend to WNA16MoEBackend.""" + mapping = { + "marlin": WNA16MoEBackend.MARLIN, + "humming": WNA16MoEBackend.HUMMING, + "flashinfer_trtllm": WNA16MoEBackend.FLASHINFER_TRTLLM, + } + if backend := mapping.get(runner_backend): + return backend + raise ValueError( + f"moe_backend='{runner_backend}' is not supported for WNA16 MoE. " + f"Expected one of {list(mapping.keys())}." + ) + + def select_wna16_moe_backend( config: FusedMoEConfig, weight_key: QuantKey, @@ -146,6 +176,14 @@ def select_wna16_moe_backend( return backend, k_cls raise ValueError(_make_log_unsupported(backend, reason)) + # Handle explicit moe_backend from user. + runner_backend = config.moe_backend + if runner_backend != "auto": + requested_backend = map_wna16_backend(runner_backend) + return _return_or_raise( + requested_backend, config, weight_key, None, activation_format + ) + # Select kernels in order of backend. AVAILABLE_BACKENDS = _get_priority_backends() @@ -210,6 +248,8 @@ def make_wna16_moe_kernel( moe_quant_config: FusedMoEQuantConfig, moe_config: FusedMoEConfig, experts_cls: type[mk.FusedMoEExperts], + backend: WNA16MoEBackend = WNA16MoEBackend.MARLIN, + layer: torch.nn.Module | None = None, is_k_full: bool = False, w13_g_idx: torch.Tensor | None = None, w2_g_idx: torch.Tensor | None = None, @@ -228,14 +268,18 @@ def make_wna16_moe_kernel( ) # Currently, we only support TrtLlmMxint4ExpertsMonolithic, MarlinExperts, - # BatchedMarlinExperts, XPUExpertsWNA16, and CPUExpertsInt4 - assert experts_cls in ( + # BatchedMarlinExperts, XPUExpertsWNA16, CPUExpertsInt4, and the Humming + # grouped/indexed experts. + allowed_experts: tuple[type[mk.FusedMoEExperts], ...] = ( MarlinExperts, BatchedMarlinExperts, TrtLlmMxint4ExpertsMonolithic, XPUExpertsWNA16, CPUExpertsInt4, ) + if backend == WNA16MoEBackend.HUMMING: + allowed_experts += tuple(backend_to_kernel_cls(WNA16MoEBackend.HUMMING)) + assert experts_cls in allowed_experts is_monolithic = experts_cls.is_monolithic() @@ -251,7 +295,10 @@ def make_wna16_moe_kernel( logger.info_once("Using %s", prepare_finalize.__class__.__name__, scope="local") extra_args: dict[str, Any] = {} - if issubclass(experts_cls, MarlinExpertsBase): + if backend == WNA16MoEBackend.HUMMING: + assert layer is not None + extra_args = {"layer": layer} + elif issubclass(experts_cls, MarlinExpertsBase): extra_args = { "w13_g_idx": w13_g_idx, "w2_g_idx": w2_g_idx, @@ -941,6 +988,35 @@ def _process_weights_xpu( ) +def _humming_wna16_weight_schema( + quant_config: QuantizationConfig | QuantizationArgs | None, +) -> dict[str, Any]: + """Humming weight schema for a WNA16 checkpoint, derived from the quant + config rather than the running kernel.""" + from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig + from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig + + if isinstance(quant_config, AutoAWQConfig): + return { + "quant_method": "awq", + "bits": quant_config.weight_bits, + "group_size": quant_config.group_size, + "zero_point": quant_config.zero_point, + } + if isinstance(quant_config, AutoGPTQConfig): + return { + "quant_method": "gptq", + "bits": quant_config.weight_bits, + "group_size": quant_config.group_size, + "desc_act": quant_config.desc_act, + "sym": quant_config.is_sym, + } + raise TypeError( + "Humming WNA16 MoE requires AutoAWQConfig or AutoGPTQConfig, " + f"got {type(quant_config).__name__}." + ) + + def convert_to_wna16_moe_kernel_format( backend: WNA16MoEBackend, layer: torch.nn.Module, @@ -956,26 +1032,30 @@ def convert_to_wna16_moe_kernel_format( w2_qzeros: torch.Tensor | None = None, w13_bias: torch.Tensor | None = None, w2_bias: torch.Tensor | None = None, -) -> tuple[ - torch.Tensor, # w13_qweight - torch.Tensor, # w2_qweight - torch.Tensor, # w13_scales - torch.Tensor, # w2_scales - torch.Tensor | None, # w13_g_idx - torch.Tensor | None, # w2_g_idx - torch.Tensor | None, # w13_g_idx_sort_indices - torch.Tensor | None, # w2_g_idx_sort_indices - torch.Tensor | None, # w13_qzeros - torch.Tensor | None, # w2_qzeros - torch.Tensor | None, # w13_input_global_scale - torch.Tensor | None, # w2_input_global_scale - torch.Tensor | None, # w13_bias - torch.Tensor | None, # w2_bias -]: +) -> ( + tuple[ + torch.Tensor, # w13_qweight + torch.Tensor, # w2_qweight + torch.Tensor, # w13_scales + torch.Tensor, # w2_scales + torch.Tensor | None, # w13_g_idx + torch.Tensor | None, # w2_g_idx + torch.Tensor | None, # w13_g_idx_sort_indices + torch.Tensor | None, # w2_g_idx_sort_indices + torch.Tensor | None, # w13_qzeros + torch.Tensor | None, # w2_qzeros + torch.Tensor | None, # w13_input_global_scale + torch.Tensor | None, # w2_input_global_scale + torch.Tensor | None, # w13_bias + torch.Tensor | None, # w2_bias + ] + | None +): """Dispatch weight post-processing to the appropriate per-backend handler. To add a new backend, implement a ``_process_weights_`` helper and - add a branch here. + add a branch here. Backends that rewrite the layer's parameters in place + (e.g. Humming) return ``None``; the caller then skips the param scatter. Args: backend: the selected ``WNA16MoEBackend``. @@ -983,6 +1063,16 @@ def convert_to_wna16_moe_kernel_format( quant_config: the ``QuantizationConfig`` for this layer. input_dtype: optional activation dtype, usually should be 16 bit. """ + if backend == WNA16MoEBackend.HUMMING: + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + convert_to_humming_moe_kernel_format, + ) + + convert_to_humming_moe_kernel_format( + layer, quant_config=_humming_wna16_weight_schema(quant_config) + ) + return None + if backend in ( WNA16MoEBackend.MARLIN, WNA16MoEBackend.BATCHED_MARLIN, diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py index 06b622a6c4b..b9086cfa48a 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py @@ -25,6 +25,7 @@ _SUPPORTED_BACKENDS = ( # is_supported_config passes (gfx950 + flydsl installed + not EP). On other # devices / no flydsl / EP it is skipped and native is used. Fp8MoeBackend.AITER_MXFP8, + Fp8MoeBackend.HUMMING, ) _BACKEND_NAME_MAP: dict[str, Fp8MoeBackend] = { @@ -34,6 +35,7 @@ _BACKEND_NAME_MAP: dict[str, Fp8MoeBackend] = { "xpu": Fp8MoeBackend.XPU, "aiter": Fp8MoeBackend.AITER_MXFP8, "triton": Fp8MoeBackend.TRITON_MXFP8, + "humming": Fp8MoeBackend.HUMMING, } diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 408c69fea09..f295163568d 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -43,6 +43,7 @@ class NvFp4MoeBackend(Enum): FLASHINFER_B12X = "FLASHINFER_B12X" VLLM_CUTLASS = "VLLM_CUTLASS" MARLIN = "MARLIN" + HUMMING = "HUMMING" EMULATION = "EMULATION" @@ -119,6 +120,18 @@ def backend_to_kernel_cls( ) return [MarlinExperts] + elif backend == NvFp4MoeBackend.HUMMING: + from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( + BatchedHummingGroupedExperts, + HummingGroupedExperts, + HummingIndexedExperts, + ) + + return [ + BatchedHummingGroupedExperts, + HummingGroupedExperts, + HummingIndexedExperts, + ] elif backend == NvFp4MoeBackend.EMULATION: from vllm.model_executor.layers.fused_moe.experts.nvfp4_emulation_moe import ( Nvfp4QuantizationEmulationTritonExperts, @@ -138,6 +151,7 @@ def map_nvfp4_backend(runner_backend: MoEBackend) -> NvFp4MoeBackend: "flashinfer_cutedsl": NvFp4MoeBackend.FLASHINFER_CUTEDSL, "flashinfer_b12x": NvFp4MoeBackend.FLASHINFER_B12X, "marlin": NvFp4MoeBackend.MARLIN, + "humming": NvFp4MoeBackend.HUMMING, "emulation": NvFp4MoeBackend.EMULATION, } if backend := mapping.get(runner_backend): @@ -169,6 +183,7 @@ def select_nvfp4_moe_backend( NvFp4MoeBackend.FLASHINFER_CUTLASS, NvFp4MoeBackend.VLLM_CUTLASS, NvFp4MoeBackend.MARLIN, + NvFp4MoeBackend.HUMMING, NvFp4MoeBackend.EMULATION, ] @@ -346,6 +361,41 @@ def convert_to_nvfp4_moe_kernel_format( a2_scale=a2_scale, is_act_and_mul=is_act_and_mul, ) + elif nvfp4_backend == NvFp4MoeBackend.HUMMING: + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + convert_to_humming_moe_kernel_format, + ) + + # Discriminate the source checkpoint layout by its global-scale param: + # compressed-tensors uses *_weight_global_scale, modelopt *_weight_scale_2. + # The logical schema is identical (nvfp4 group-16); only the on-layer + # param names differ. TODO: normalize both methods to a single canonical + # layout upstream so the oracle needs neither the probe nor the re-alias. + if hasattr(layer, "w13_weight_global_scale"): + quant_config = { + "quant_method": "compressed-tensors", + "format": "nvfp4-pack-quantized", + "type": "float", + "num_bits": 4, + "strategy": "group", + "group_size": 16, + } + # CT pack-quantized reads `weight_packed`; the method renamed it to + # `weight`. Re-alias (convert replaces all params anyway). + layer.w13_weight_packed = layer.w13_weight + layer.w2_weight_packed = layer.w2_weight + else: + quant_config = {"quant_method": "modelopt", "quant_algo": "nvfp4"} + + convert_to_humming_moe_kernel_format(layer, quant_config=quant_config) + a13_scale = None + a2_scale = None + w13 = layer.w13_weight + w13_scale = layer.w13_weight_scale + w13_scale_2 = getattr(layer, "w13_global_scale", None) + w2 = layer.w2_weight + w2_scale = layer.w2_weight_scale + w2_scale_2 = getattr(layer, "w2_global_scale", None) elif nvfp4_backend == NvFp4MoeBackend.MARLIN: a13_scale = None a2_scale = None @@ -418,8 +468,17 @@ def make_nvfp4_moe_quant_config( a13_scale: torch.Tensor, a2_scale: torch.Tensor, swiglu_limit: float | None = None, + layer: torch.nn.Module | None = None, ) -> FusedMoEQuantConfig: - if backend == NvFp4MoeBackend.MARLIN: + if backend == NvFp4MoeBackend.HUMMING: + from vllm.model_executor.layers.fused_moe import RoutedExperts + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + get_humming_moe_quant_config, + ) + + assert isinstance(layer, RoutedExperts) + return get_humming_moe_quant_config(layer) + elif backend == NvFp4MoeBackend.MARLIN: return nvfp4_w4a16_moe_quant_config( g1_alphas=w13_scale_2, g2_alphas=w2_scale_2, @@ -467,7 +526,9 @@ def make_nvfp4_moe_kernel( moe_quant_config: FusedMoEQuantConfig, moe_config: FusedMoEConfig, experts_cls: type[mk.FusedMoEExperts], + backend: NvFp4MoeBackend, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, + layer: torch.nn.Module | None = None, ) -> mk.FusedMoEKernel: # Create Prepare/Finalize. prepare_finalize = maybe_make_prepare_finalize( @@ -481,6 +542,11 @@ def make_nvfp4_moe_kernel( logger.info_once("Using %s", prepare_finalize.__class__.__name__) + extra_kwargs = {} + if backend == NvFp4MoeBackend.HUMMING: + assert layer is not None + extra_kwargs = {"layer": layer} + # Create Experts. if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts: max_num_tokens = prepare_finalize.max_num_tokens_per_rank() @@ -490,11 +556,13 @@ def make_nvfp4_moe_kernel( quant_config=moe_quant_config, max_num_tokens=max_num_tokens, num_dispatchers=prepare_finalize.num_dispatchers(), + **extra_kwargs, ) else: experts = experts_cls( moe_config=moe_config, quant_config=moe_quant_config, + **extra_kwargs, ) kernel = mk.FusedMoEKernel( diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index cda0eaf7300..639fce5234c 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -232,7 +232,9 @@ def select_unquantized_moe_backend( raise ValueError(_make_log_unsupported(backend, reason)) runner_backend = moe_config.moe_backend - if runner_backend != "auto": + # 'humming' is quantization-only; an unquantized layer (e.g. excluded via + # modules_to_not_convert) falls through to auto instead of erroring. + if runner_backend not in ["auto", "humming"]: requested_backend = map_unquantized_backend(runner_backend) if ( activation_format == mk.FusedMoEActivationFormat.BatchedExperts diff --git a/vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py b/vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py index e4f4d497568..c66e43ccba8 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py @@ -270,7 +270,7 @@ def convert_to_w4a8_int8_moe_format( """ # Derive dimensions from tensor shapes E = w13_weight.shape[0] # num_experts - I2 = w13_weight.shape[1] # w13_out_features (2*IN) + w13_out_features = w13_weight.shape[1] # 2 * intermediate_size H = w13_weight.shape[2] # w13_in_features (hidden_size) IN = w2_weight.shape[2] # w2_in_features (intermediate_size) w2_out_features = w2_weight.shape[1] # Should equal H @@ -286,7 +286,7 @@ def convert_to_w4a8_int8_moe_format( w13_weight_scale[e], # [2I, H/g or 1] w13_bias[e] if w13_bias is not None else None, # [2I] H, - I2, + w13_out_features, group_size, ) ) diff --git a/vllm/model_executor/layers/hpc/rope_norm.py b/vllm/model_executor/layers/hpc/rope_norm.py index 7eee2a6eb30..2b07d949c02 100644 --- a/vllm/model_executor/layers/hpc/rope_norm.py +++ b/vllm/model_executor/layers/hpc/rope_norm.py @@ -7,6 +7,7 @@ Decoupled from HpcAttentionImpl; extra params are passed via layer attrs. from __future__ import annotations +import importlib.util from enum import IntEnum from typing import Any @@ -122,6 +123,12 @@ class HpcRopeNorm(CustomOp, HpcModule): qk_norm_policy: QkNormPolicy = QkNormPolicy.ROPE_THEN_NORM, ) -> None: super().__init__() + if importlib.util.find_spec("hpc") is None: + raise ImportError( + "HPCRopeNorm requires the hpc module to be installed. " + "Please install it from https://github.com/Tencent/hpc-ops" + ) + self.num_heads = num_heads self.num_kv_heads = num_kv_heads self.head_dim = head_dim @@ -172,6 +179,15 @@ class HpcRopeNorm(CustomOp, HpcModule): self.layer_name: str | None = None self.register_layer_name(layer_name) + import hpc + + if self.use_fp8: + self._quant_type = ( + hpc.QuantType.QPERTOKEN_PERHEAD_KPERTENSOR_VPERTENSOR.value + ) + else: + self._quant_type = None + @classmethod def support( cls, @@ -304,12 +320,6 @@ class HpcRopeNorm(CustomOp, HpcModule): self.knorm_weight if self.qk_norm_policy != QkNormPolicy.NONE else None ) - # Dynamic per-token-per-head Q quant + per-tensor K/V (dqskv). - # rope_norm_store_kv_fp8 is registered as a torch op whose ``quant_policy`` - # argument is typed as ``int``; pybind cannot cast the hpc.QuantType enum - # automatically, so pass its integer ``.value``. - QUANT_POLICY_DQSKV = hpc.QuantType.QPERTOKEN_PERHEAD_KPERTENSOR_VPERTENSOR.value - # --- Prefill --- if num_prefill_reqs > 0: seq_lens_prefill = attn_metadata.seq_lens[num_decode_reqs:] @@ -333,7 +343,7 @@ class HpcRopeNorm(CustomOp, HpcModule): is_prefill=True, k_scale=k_scale, v_scale=v_scale, - quant_policy=QUANT_POLICY_DQSKV, + quant_policy=self._quant_type, max_seqlens=max_seqlens, q_norm_weight=q_norm_weight, k_norm_weight=k_norm_weight, @@ -364,9 +374,7 @@ class HpcRopeNorm(CustomOp, HpcModule): qkv_decode = qkv[:num_decode_tokens] # Single-token decode: q_index is the per-request prefix sum # [0, 1, ..., num_decode_reqs]. - qo_indptr_decode = torch.arange( - num_decode_reqs + 1, dtype=torch.int32, device=qkv.device - ) + decode_query_len = attn_metadata.decode_query_len out_q_decode = output[:num_decode_tokens] if self.use_fp8: @@ -376,13 +384,13 @@ class HpcRopeNorm(CustomOp, HpcModule): qkv=qkv_decode, cos_sin=self.cos_sin_cache, num_seqlen_per_req=num_seq_kvcache, - q_index=qo_indptr_decode, + q_index=attn_metadata.qo_indptr_decode, kvcache_indices=block_table_decode, is_prefill=False, k_scale=k_scale, v_scale=v_scale, - quant_policy=QUANT_POLICY_DQSKV, - max_seqlens=1, + quant_policy=self._quant_type, + max_seqlens=decode_query_len, q_norm_weight=q_norm_weight, k_norm_weight=k_norm_weight, qk_norm_policy=self.qk_norm_policy, @@ -398,7 +406,7 @@ class HpcRopeNorm(CustomOp, HpcModule): qkv_decode, self.cos_sin_cache, num_seq_kvcache, - qo_indptr_decode, + attn_metadata.qo_indptr_decode, block_table_decode, False, # is_prefill q_norm_weight=q_norm_weight, diff --git a/vllm/model_executor/layers/mamba/linear/bailing_linear_attn.py b/vllm/model_executor/layers/mamba/linear/bailing_linear_attn.py index dd963f829d8..a00fbc74bf8 100644 --- a/vllm/model_executor/layers/mamba/linear/bailing_linear_attn.py +++ b/vllm/model_executor/layers/mamba/linear/bailing_linear_attn.py @@ -30,11 +30,10 @@ from vllm.model_executor.layers.mamba.linear.base import LinearAttention from vllm.model_executor.layers.mamba.linear.minimax_linear_attn import ( MiniMaxText01LinearAttention, MiniMaxText01LinearKernel, - clear_linear_attention_cache_for_new_sequences, linear_attention_decode, - linear_attention_prefill_and_mix, ) from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.linear_attn import LinearAttentionMetadata @@ -57,6 +56,290 @@ def _build_rope_parameters(config: PretrainedConfig) -> dict | None: return rope_parameters or None +def clear_linear_attention_cache_for_new_sequences( + kv_cache: torch.Tensor, + state_indices_tensor: torch.Tensor, + attn_metadata: LinearAttentionMetadata, +) -> None: + num_prefills = getattr(attn_metadata, "num_prefills", 0) + if num_prefills <= 0: + return + + num_decodes = getattr(attn_metadata, "num_decodes", 0) + prefill_state_indices = getattr(attn_metadata, "state_indices_tensor_p", None) + for prefill_idx in range(num_prefills): + if num_decodes + prefill_idx + 1 >= len(attn_metadata.query_start_loc): + break + q_start = attn_metadata.query_start_loc[num_decodes + prefill_idx] + q_end = attn_metadata.query_start_loc[num_decodes + prefill_idx + 1] + query_len = q_end - q_start + context_len = attn_metadata.seq_lens[num_decodes + prefill_idx] - query_len + if context_len == 0: + if prefill_state_indices is not None: + block_to_clear = prefill_state_indices[prefill_idx] + else: + block_to_clear = state_indices_tensor[num_decodes + prefill_idx] + kv_cache[block_to_clear, ...] = 0 + + +def linear_attention_prefill_and_mix( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kv_cache: torch.Tensor, + state_indices_tensor: torch.Tensor, + attn_metadata: LinearAttentionMetadata, + slope_rate: torch.Tensor, + block_size: int, + decode_fn, + prefix_fn, + layer_idx: int | None = None, +) -> torch.Tensor: + hidden = [] + req_offset = getattr(attn_metadata, "num_decodes", 0) + prefill_state_indices = getattr(attn_metadata, "state_indices_tensor_p", None) + for _prefill_idx in range(getattr(attn_metadata, "num_prefills", 0)): + if req_offset + _prefill_idx + 1 >= len(attn_metadata.query_start_loc): + break + if prefill_state_indices is not None and _prefill_idx >= len( + prefill_state_indices + ): + break + if prefill_state_indices is None and _prefill_idx >= len(state_indices_tensor): + break + _start = attn_metadata.query_start_loc[req_offset + _prefill_idx] + _end = attn_metadata.query_start_loc[req_offset + _prefill_idx + 1] + if prefill_state_indices is not None: + slot_id = prefill_state_indices[_prefill_idx] + else: + slot_id = state_indices_tensor[req_offset + _prefill_idx] + qs = q[_start:_end].transpose(0, 1).contiguous() + ks = k[_start:_end].transpose(0, 1).contiguous() + vs = v[_start:_end].transpose(0, 1).contiguous() + slice_layer_cache = kv_cache[slot_id, ...] + out_slice = prefix_fn( + qs, + ks, + vs, + slice_layer_cache, + slope_rate, + block_size, + layer_idx=layer_idx, + ) + hidden.append(out_slice.contiguous()) + + if attn_metadata.num_decode_tokens > 0: + hidden_decode = decode_fn( + q, k, v, kv_cache, state_indices_tensor, attn_metadata + ) + hidden.insert(0, hidden_decode) + + if not hidden: + return torch.empty((0, q.size(1) * q.size(2)), device=q.device, dtype=q.dtype) + + hidden = torch.concat(hidden, dim=0).contiguous() + return hidden + + +@triton.jit +def _bailing_linear_attn_decode_spec_step_kernel( + q_ptr, + k_ptr, + v_ptr, + kv_cache_ptr, + slope_rate_ptr, + state_indices_ptr, + query_start_loc_ptr, + num_accepted_tokens_ptr, + output_ptr, + q_start: tl.constexpr, + D: tl.constexpr, + q_b_stride, + q_h_stride, + q_d_stride, + k_b_stride, + k_h_stride, + k_d_stride, + v_b_stride, + v_h_stride, + v_d_stride, + cache_b_stride, + cache_h_stride, + cache_d0_stride, + cache_d1_stride, + state_indices_b_stride, + state_indices_t_stride, + output_b_stride, + output_d_stride, + DRAFT_IDX: tl.constexpr, + STATE_WIDTH: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + req_id = tl.program_id(0) + head_id = tl.program_id(1) + block_id = tl.program_id(2) + + req_start = tl.load(query_start_loc_ptr + req_id).to(tl.int64) + req_end = tl.load(query_start_loc_ptr + req_id + 1).to(tl.int64) + query_len = req_end - req_start + if query_len <= DRAFT_IDX: + return + + dst_slot = tl.load( + state_indices_ptr + + req_id * state_indices_b_stride + + DRAFT_IDX * state_indices_t_stride + ).to(tl.int64) + if dst_slot == -1: + return + + if DRAFT_IDX == 0: + accepted_offset = tl.load(num_accepted_tokens_ptr + req_id).to(tl.int64) - 1 + accepted_offset = tl.maximum(accepted_offset, 0) + accepted_offset = tl.minimum(accepted_offset, STATE_WIDTH - 1) + src_slot = tl.load( + state_indices_ptr + + req_id * state_indices_b_stride + + accepted_offset * state_indices_t_stride + ).to(tl.int64) + else: + src_slot = tl.load( + state_indices_ptr + + req_id * state_indices_b_stride + + (DRAFT_IDX - 1) * state_indices_t_stride + ).to(tl.int64) + if src_slot == -1: + return + + token_idx = req_start - q_start + DRAFT_IDX + qk_offsets = tl.arange(0, D) + v_offsets = tl.arange(0, BLOCK_SIZE) + block_id * BLOCK_SIZE + qk_mask = qk_offsets < D + v_mask = v_offsets < D + kv_mask = qk_mask[:, None] & v_mask[None, :] + + q = tl.load( + q_ptr + token_idx * q_b_stride + head_id * q_h_stride + qk_offsets * q_d_stride, + mask=qk_mask, + other=0.0, + ) + k = tl.load( + k_ptr + token_idx * k_b_stride + head_id * k_h_stride + qk_offsets * k_d_stride, + mask=qk_mask, + other=0.0, + ) + v = tl.load( + v_ptr + token_idx * v_b_stride + head_id * v_h_stride + v_offsets * v_d_stride, + mask=v_mask, + other=0.0, + ) + + cache_offsets = ( + qk_offsets[:, None] * cache_d0_stride + v_offsets[None, :] * cache_d1_stride + ) + src_cache_ptr = ( + kv_cache_ptr + + src_slot * cache_b_stride + + head_id * cache_h_stride + + cache_offsets + ) + dst_cache_ptr = ( + kv_cache_ptr + + dst_slot * cache_b_stride + + head_id * cache_h_stride + + cache_offsets + ) + + slope = tl.load(slope_rate_ptr + head_id) + decay = tl.exp(-slope) + kv_old = tl.load(src_cache_ptr, mask=kv_mask, other=0.0) + kv_new = k[:, None] * v[None, :] + decay * kv_old + + output = tl.sum(q[:, None].to(tl.float32) * kv_new, axis=0) + tl.store(dst_cache_ptr, kv_new, mask=kv_mask) + tl.store( + output_ptr + + token_idx * output_b_stride + + (head_id * D + v_offsets) * output_d_stride, + output, + mask=v_mask, + ) + + +def bailing_linear_attention_decode_spec( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kv_cache: torch.Tensor, + slope_rate: torch.Tensor, + state_indices_tensor: torch.Tensor, + query_start_loc: torch.Tensor, + num_accepted_tokens: torch.Tensor, + q_start: int, + q_end: int | None, + slot_start: int, + slot_end: int | None, + block_size: int, +) -> torch.Tensor: + q_decode = q[q_start:q_end] + k_decode = k[q_start:q_end] + v_decode = v[q_start:q_end] + hidden = torch.empty( + (q_decode.shape[0], q.shape[1] * q.shape[2]), + device=q.device, + dtype=q.dtype, + ) + hidden.zero_() + + state_indices_tensor = state_indices_tensor[slot_start:slot_end] + query_start_loc = query_start_loc.to(device=q.device) + + batch_size = state_indices_tensor.shape[0] + num_heads = q_decode.shape[1] + head_dim = q_decode.shape[2] + assert k_decode.shape == (q_decode.shape[0], num_heads, head_dim) + assert v_decode.shape == (q_decode.shape[0], num_heads, head_dim) + state_width = state_indices_tensor.shape[1] + + grid = (batch_size, num_heads, triton.cdiv(head_dim, block_size)) + for draft_idx in range(state_width): + _bailing_linear_attn_decode_spec_step_kernel[grid]( + q_decode, + k_decode, + v_decode, + kv_cache, + slope_rate, + state_indices_tensor, + query_start_loc, + num_accepted_tokens[:batch_size], + hidden, + q_start, + head_dim, + q_decode.stride(0), + q_decode.stride(1), + q_decode.stride(2), + k_decode.stride(0), + k_decode.stride(1), + k_decode.stride(2), + v_decode.stride(0), + v_decode.stride(1), + v_decode.stride(2), + kv_cache.stride(0), + kv_cache.stride(1), + kv_cache.stride(2), + kv_cache.stride(3), + state_indices_tensor.stride(0), + state_indices_tensor.stride(1), + hidden.stride(0), + hidden.stride(1), + DRAFT_IDX=draft_idx, + STATE_WIDTH=state_width, + BLOCK_SIZE=block_size, + ) + + return hidden + + class BailingGroupRMSNormGate(RMSNormGated): def __init__( self, @@ -208,6 +491,13 @@ class BailingMoELinearAttention(LinearAttention): raise ValueError(f"Duplicate layer name: {prefix}") compilation_config.static_forward_context[prefix] = self + def get_attn_backend(self): + from vllm.v1.attention.backends.linear_attn import ( + BailingLinearAttentionBackend, + ) + + return BailingLinearAttentionBackend + @staticmethod def weight_direct_load(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: """Load weight for linear attention layers. @@ -368,13 +658,42 @@ class BailingMoELinearAttention(LinearAttention): def _decode_infer(self, q, k, v, kv_cache, state_indices_tensor, attn_metadata): """Handle decode (single token per sequence).""" + decode_state_indices = getattr(attn_metadata, "state_indices_tensor_d", None) + num_accepted_tokens = getattr(attn_metadata, "num_accepted_tokens", None) + query_start_loc = getattr(attn_metadata, "query_start_loc_d", None) + if ( + decode_state_indices is not None + and decode_state_indices.dim() > 1 + and num_accepted_tokens is not None + and query_start_loc is not None + ): + return bailing_linear_attention_decode_spec( + q, + k, + v, + kv_cache, + self.tp_slope, + decode_state_indices, + query_start_loc, + num_accepted_tokens, + q_start=0, + q_end=attn_metadata.num_decode_tokens, + slot_start=0, + slot_end=attn_metadata.num_decodes, + block_size=32, + ) + decode_state_indices = ( + state_indices_tensor + if decode_state_indices is None + else decode_state_indices + ) hidden = linear_attention_decode( q, k, v, kv_cache, self.tp_slope, - state_indices_tensor, + decode_state_indices, q_start=0, q_end=attn_metadata.num_decode_tokens, slot_start=0, diff --git a/vllm/model_executor/layers/mamba/short_conv.py b/vllm/model_executor/layers/mamba/short_conv.py index 79976dfff14..e7e36f2fc53 100644 --- a/vllm/model_executor/layers/mamba/short_conv.py +++ b/vllm/model_executor/layers/mamba/short_conv.py @@ -23,6 +23,7 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( causal_conv1d_fn, causal_conv1d_update, ) +from vllm.platforms import current_platform from vllm.utils.torch_utils import direct_register_custom_op from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum @@ -90,7 +91,94 @@ class ShortConv(MambaBase, CustomOp): hidden_states: torch.Tensor, output: torch.Tensor, ): - return + # Reference torch causal conv1d; runs on all CPU platforms. AMX kernels + # for causal conv can be plugged in here later. + from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( + causal_conv1d_torch, + causal_conv1d_update_torch, + ) + + forward_context = get_forward_context() + attn_metadata_raw = forward_context.attn_metadata + attn_metadata: AttentionMetadata | None = None + if attn_metadata_raw is not None: + assert isinstance(attn_metadata_raw, dict) + attn_metadata = attn_metadata_raw[self.prefix] + assert isinstance(attn_metadata, ShortConvAttentionMetadata) + + BCx, _ = self.in_proj(hidden_states) + B, C, x = BCx.chunk(3, dim=-1) + + # (dim, kernel_size) — same reshape as forward_cuda + conv_weights = self.conv.weight.view( + self.conv.weight.size(0), self.conv.weight.size(2) + ) + + if attn_metadata is None: + # Profile run — output value doesn't matter + Bx = (B * x).contiguous() + output_tensor, _ = self.out_proj(C * Bx) + output[: hidden_states.shape[0]] = output_tensor + return + + conv_state = ( + self.kv_cache[0] + if is_conv_state_dim_first() + else self.kv_cache[0].transpose(-1, -2) + ) # (num_blocks, dim, state_len) + + num_prefills = attn_metadata.num_prefills + num_decodes = attn_metadata.num_decode_tokens + num_prefill_tokens = attn_metadata.num_prefill_tokens + has_prefill = num_prefills > 0 + has_decode = num_decodes > 0 + num_actual_tokens = num_decodes + num_prefill_tokens + + B_d, B_p = torch.split( + B[:num_actual_tokens], [num_decodes, num_prefill_tokens], dim=0 + ) + C_d, C_p = torch.split( + C[:num_actual_tokens], [num_decodes, num_prefill_tokens], dim=0 + ) + x_d, x_p = torch.split( + x[:num_actual_tokens], [num_decodes, num_prefill_tokens], dim=0 + ) + + conv_output_list = [] + + if has_prefill: + assert attn_metadata.state_indices_tensor_p is not None + Bx_p = (B_p * x_p).transpose(0, 1) # (dim, num_prefill_tokens) + out_p = causal_conv1d_torch( + Bx_p, + conv_weights, + self.conv.bias, + conv_state, + attn_metadata.query_start_loc_p, + attn_metadata.state_indices_tensor_p.flatten(), + attn_metadata.has_initial_states_p, + activation=None, + ).transpose(0, 1)[:num_prefill_tokens] # (num_prefill_tokens, dim) + conv_output_list.append(C_p * out_p) + + if has_decode: + assert attn_metadata.state_indices_tensor_d is not None + state_indices_d = attn_metadata.state_indices_tensor_d.flatten() + Bx_d = (B_d * x_d).unsqueeze(-1) # (num_decodes, dim, 1) + # Advanced indexing returns a copy; update in-place then scatter back + gathered = conv_state[state_indices_d] # (num_decodes, dim, state_len) + out_d = causal_conv1d_update_torch( + Bx_d, + gathered, + conv_weights, + self.conv.bias, + activation=None, + ).squeeze(-1) # (num_decodes, dim) + conv_state[state_indices_d] = gathered + conv_output_list.insert(0, C_d * out_d) + + hidden_states_out = torch.vstack(conv_output_list) + output[:num_actual_tokens], _ = self.out_proj(hidden_states_out) def forward( self, @@ -235,7 +323,10 @@ def short_conv( ) -> None: forward_context: ForwardContext = get_forward_context() self = forward_context.no_compile_layers[layer_name] - self.forward_cuda(hidden_states=hidden_states, output=output) + if not current_platform.is_cpu(): + self.forward_cuda(hidden_states=hidden_states, output=output) + else: + self.forward_native(hidden_states=hidden_states, output=output) def short_conv_fake( diff --git a/vllm/model_executor/layers/quantization/auto_awq.py b/vllm/model_executor/layers/quantization/auto_awq.py index cebfad7e596..58104fa7d25 100644 --- a/vllm/model_executor/layers/quantization/auto_awq.py +++ b/vllm/model_executor/layers/quantization/auto_awq.py @@ -26,6 +26,7 @@ from vllm.model_executor.layers.fused_moe import ( UnquantizedFusedMoEMethod, ) from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import ( + WNA16MoEBackend, convert_to_wna16_moe_kernel_format, make_wna16_moe_kernel, make_wna16_moe_quant_config, @@ -663,6 +664,26 @@ class AutoAWQMoEMethod(FusedMoEMethodBase): layer.workspace = marlin_make_workspace_new(device, 4) def process_weights_after_loading(self, layer: RoutedExperts) -> None: + converted = convert_to_wna16_moe_kernel_format( + backend=self.wna16_moe_backend, + layer=layer, + quant_config=self.quant_config, + input_dtype=self.input_dtype, + w13=layer.w13_qweight, + w2=layer.w2_qweight, + w13_scale=layer.w13_scales, + w2_scale=layer.w2_scales, + w13_qzeros=layer.w13_qzeros, + w2_qzeros=layer.w2_qzeros, + w13_bias=getattr(layer, "w13_bias", None), + w2_bias=getattr(layer, "w2_bias", None), + ) + + if converted is None: + # Backend rewrote the layer's params in place (e.g. Humming). + self._setup_kernel(layer) + return + ( w13, w2, @@ -678,20 +699,7 @@ class AutoAWQMoEMethod(FusedMoEMethodBase): w2_input_global_scale, w13_bias, w2_bias, - ) = convert_to_wna16_moe_kernel_format( - backend=self.wna16_moe_backend, - layer=layer, - quant_config=self.quant_config, - input_dtype=self.input_dtype, - w13=layer.w13_qweight, - w2=layer.w2_qweight, - w13_scale=layer.w13_scales, - w2_scale=layer.w2_scales, - w13_qzeros=layer.w13_qzeros, - w2_qzeros=layer.w2_qzeros, - w13_bias=getattr(layer, "w13_bias", None), - w2_bias=getattr(layer, "w2_bias", None), - ) + ) = converted replace_parameter(layer, "w13_qweight", w13) replace_parameter(layer, "w2_qweight", w2) @@ -734,6 +742,8 @@ class AutoAWQMoEMethod(FusedMoEMethodBase): moe_quant_config=self.moe_quant_config, moe_config=self.moe, experts_cls=self.experts_cls, + backend=self.wna16_moe_backend, + layer=layer, is_k_full=self.is_k_full, w13_g_idx=getattr(layer, "w13_g_idx", None), w2_g_idx=getattr(layer, "w2_g_idx", None), @@ -743,6 +753,12 @@ class AutoAWQMoEMethod(FusedMoEMethodBase): ) def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig: + if self.wna16_moe_backend == WNA16MoEBackend.HUMMING: + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + get_humming_moe_quant_config, + ) + + return get_humming_moe_quant_config(layer) return make_wna16_moe_quant_config( w1_scale=layer.w13_scales, w2_scale=layer.w2_scales, @@ -783,8 +799,8 @@ class AutoAWQMoEMethod(FusedMoEMethodBase): assert self.moe_kernel is not None return self.moe_kernel.apply( hidden_states=x, - w1=layer.w13_qweight, - w2=layer.w2_qweight, + w1=layer.w13_weight, + w2=layer.w2_weight, topk_weights=topk_weights, topk_ids=topk_ids, activation=layer.activation, @@ -806,8 +822,8 @@ class AutoAWQMoEMethod(FusedMoEMethodBase): assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( hidden_states=x, - w1=layer.w13_qweight, - w2=layer.w2_qweight, + w1=layer.w13_weight, + w2=layer.w2_weight, router_logits=router_logits, activation=layer.activation, global_num_experts=layer.global_num_experts, diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py index aca76162f9a..b056f5222af 100644 --- a/vllm/model_executor/layers/quantization/auto_gptq.py +++ b/vllm/model_executor/layers/quantization/auto_gptq.py @@ -658,6 +658,26 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): "W8A8-INT8 is not supported by marlin kernel." ) + converted = convert_to_wna16_moe_kernel_format( + backend=self.wna16_moe_backend, + layer=layer, + quant_config=self.quant_config, + input_dtype=self.input_dtype, + w13=layer.w13_qweight, + w2=layer.w2_qweight, + w13_scale=layer.w13_scales, + w2_scale=layer.w2_scales, + w13_g_idx=layer.w13_g_idx, + w2_g_idx=layer.w2_g_idx, + w13_bias=getattr(layer, "w13_bias", None), + w2_bias=getattr(layer, "w2_bias", None), + ) + + if converted is None: + # Backend rewrote the layer's params in place (e.g. Humming). + self._setup_kernel(layer) + return + ( w13, w2, @@ -673,20 +693,7 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): w2_input_global_scale, w13_bias, w2_bias, - ) = convert_to_wna16_moe_kernel_format( - backend=self.wna16_moe_backend, - layer=layer, - quant_config=self.quant_config, - input_dtype=self.input_dtype, - w13=layer.w13_qweight, - w2=layer.w2_qweight, - w13_scale=layer.w13_scales, - w2_scale=layer.w2_scales, - w13_g_idx=layer.w13_g_idx, - w2_g_idx=layer.w2_g_idx, - w13_bias=getattr(layer, "w13_bias", None), - w2_bias=getattr(layer, "w2_bias", None), - ) + ) = converted replace_parameter(layer, "w13_qweight", w13) replace_parameter(layer, "w2_qweight", w2) @@ -733,6 +740,10 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): "w2_bias", torch.nn.Parameter(w2_bias, requires_grad=False) ) + # The modular kernel reads w13_weight/w2_weight; marlin keeps *_qweight. + layer.w13_weight = layer.w13_qweight + layer.w2_weight = layer.w2_qweight + self._setup_kernel(layer) def _setup_kernel(self, layer: RoutedExperts) -> None: @@ -743,15 +754,24 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): moe_quant_config=self.moe_quant_config, moe_config=self.moe, experts_cls=self.experts_cls, + backend=self.wna16_moe_backend, + layer=layer, is_k_full=self.is_k_full, - w13_g_idx=layer.w13_g_idx, - w2_g_idx=layer.w2_g_idx, + w13_g_idx=getattr(layer, "w13_g_idx", None), + w2_g_idx=getattr(layer, "w2_g_idx", None), w13_g_idx_sort_indices=getattr(layer, "w13_g_idx_sort_indices", None), w2_g_idx_sort_indices=getattr(layer, "w2_g_idx_sort_indices", None), routing_tables=layer._expert_routing_tables(), ) def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig: + if self.wna16_moe_backend == WNA16MoEBackend.HUMMING: + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + get_humming_moe_quant_config, + ) + + return get_humming_moe_quant_config(layer) + from vllm.model_executor.layers.fused_moe.config import ( gptq_marlin_moe_quant_config, ) @@ -795,8 +815,8 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): assert self.moe_kernel is not None return self.moe_kernel.apply( hidden_states=x, - w1=layer.w13_qweight, - w2=layer.w2_qweight, + w1=layer.w13_weight, + w2=layer.w2_weight, topk_weights=topk_weights, topk_ids=topk_ids, activation=layer.activation, @@ -818,8 +838,8 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( hidden_states=x, - w1=layer.w13_qweight, - w2=layer.w2_qweight, + w1=layer.w13_weight, + w2=layer.w2_weight, router_logits=router_logits, activation=layer.activation, global_num_experts=layer.global_num_experts, 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 9a051c038f9..f74ed2d7b38 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 @@ -236,7 +236,9 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): moe_quant_config=self.moe_quant_config, moe_config=self.moe, experts_cls=self.experts_cls, + backend=self.nvfp4_backend, routing_tables=layer._expert_routing_tables(), + layer=layer, ) self.moe_kernel.fused_experts.process_weights_after_loading(layer) @@ -259,6 +261,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): a13_scale=layer.w13_input_scale, a2_scale=layer.w2_input_scale, swiglu_limit=getattr(layer, "swiglu_limit", None), + layer=layer, ) def apply_monolithic( 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 a64104d3ffd..2d2190e614e 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 @@ -5,7 +5,6 @@ import torch from compressed_tensors.quantization import ( QuantizationArgs, - QuantizationStrategy, ) from vllm.logger import init_logger @@ -32,7 +31,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( ScaleDesc, ) from vllm.model_executor.utils import replace_parameter, set_weight_attrs -from vllm.platforms import CpuArchEnum, current_platform logger = init_logger(__name__) @@ -64,32 +62,17 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod): weight_quant.group_size if (weight_quant.group_size is not None) else -1 ) - # Validate scheme: weights=W4 (channel or group), - # activations=dynamic TOKEN (A8) - - # Must be dynamic per-token activations - if ( - input_quant.strategy != QuantizationStrategy.TOKEN - or not input_quant.dynamic + # make sure group size is valid + if self.group_size != -1 and ( + moe.hidden_dim % self.group_size != 0 + or moe.intermediate_size_per_partition % self.group_size != 0 ): raise ValueError( - "W4A8-int MoE needs dynamic per-token activation quantization." + f"Group size ({self.group_size}) must evenly divide both " + f"hidden size ({moe.hidden_dim}) and intermediate size per " + f"partition ({moe.intermediate_size_per_partition})." ) - if weight_quant.num_bits != 4: - raise ValueError("This method only supports 4-bit weights (num_bits=4).") - - # Arm: check _dyn ops availability - if current_platform.get_cpu_architecture() == CpuArchEnum.ARM: - try: - _ = torch.ops.aten._dyn_quant_matmul_4bit - _ = torch.ops.aten._dyn_quant_pack_4bit_weight - except AttributeError as err: - raise RuntimeError( - f"""PyTorch {torch.__version__} lacks _dyn_quant_* 4bit ops; - install a newer build.""" - ) from err - # Construct QuantKey for weights from QuantizationArgs # W4A8 INT4: 4-bit weights (stored as int8), static quantization if self.group_size == -1: 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 14ef8bf614c..35d2a378424 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 @@ -338,6 +338,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod): fp8_backend=self.fp8_backend, experts_cls=self.experts_cls, routing_tables=layer._expert_routing_tables(), + layer=layer, ) def maybe_make_prepare_finalize( @@ -355,12 +356,13 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod): fp8_backend=self.fp8_backend, w1_scale=layer.w13_weight_scale, w2_scale=layer.w2_weight_scale, - a1_scale=layer.w13_input_scale, - a2_scale=layer.w2_input_scale, + a1_scale=getattr(layer, "w13_input_scale", None), + a2_scale=getattr(layer, "w2_input_scale", None), per_act_token_quant=is_per_token, per_out_ch_quant=is_per_token, block_shape=self.weight_block_size, swiglu_limit=getattr(layer, "swiglu_limit", None), + layer=layer, ) def apply_monolithic( 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 c29472cfc6b..d304ca56bf1 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 @@ -146,6 +146,8 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod): int8_backend=self.int8_backend, w13=layer.w13_weight, w2=layer.w2_weight, + layer=layer, + w13_scale=layer.w13_weight_scale, ) replace_parameter(layer, "w13_weight", w13) replace_parameter(layer, "w2_weight", w2) @@ -153,10 +155,12 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod): self.moe_quant_config = self.get_fused_moe_quant_config(layer) assert self.experts_cls is not None self.moe_kernel = make_int8_moe_kernel( + int8_backend=self.int8_backend, moe_quant_config=self.moe_quant_config, moe_config=self.moe, experts_cls=self.experts_cls, routing_tables=layer._expert_routing_tables(), + layer=layer, ) def maybe_make_prepare_finalize( @@ -170,11 +174,13 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod): def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig: return make_int8_moe_quant_config( + int8_backend=self.int8_backend, w1_scale=layer.w13_weight_scale, w2_scale=layer.w2_weight_scale, a1_scale=layer.w13_input_scale, a2_scale=layer.w2_input_scale, per_act_token_quant=True, + layer=layer, ) def apply( 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 2e6e01ca766..468aed29a9e 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 @@ -140,6 +140,7 @@ class CompressedTensorsW8A8Mxfp8MoEMethod(CompressedTensorsMoEMethod): fp8_backend=self.fp8_backend, experts_cls=self.experts_cls, routing_tables=layer._expert_routing_tables(), + layer=layer, ) def get_fused_moe_quant_config( @@ -155,6 +156,7 @@ class CompressedTensorsW8A8Mxfp8MoEMethod(CompressedTensorsMoEMethod): swiglu_limit=getattr(layer, "swiglu_limit", None), gemm1_alpha=getattr(layer, "swiglu_alpha", None), gemm1_beta=getattr(layer, "swiglu_beta", None), + layer=layer, ) def maybe_make_prepare_finalize( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py index 46fa36180d9..8af36bcb102 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py @@ -416,6 +416,27 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # Process weights using the shared oracle infrastructure is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM + converted = convert_to_wna16_moe_kernel_format( + backend=self.wna16_backend, + layer=layer, + quant_config=self.weight_quant, + input_dtype=self.marlin_input_dtype, + w13=layer.w13_weight_packed, + w2=layer.w2_weight_packed, + w13_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + w13_g_idx=layer.w13_weight_g_idx, + w2_g_idx=layer.w2_weight_g_idx, + w13_qzeros=getattr(layer, "w13_weight_zero_point", None), + w2_qzeros=getattr(layer, "w2_weight_zero_point", None), + ) + if converted is None: + # In-place backends (e.g. Humming) are not wired through this + # marlin-only method; fail clearly rather than unpacking None. + raise NotImplementedError( + f"{type(self).__name__} does not support the " + f"{self.wna16_backend.value} MoE backend." + ) ( w13_qweight, w2_qweight, @@ -431,20 +452,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): w2_input_global_scale, _, # w13_bias _, # w2_bias - ) = convert_to_wna16_moe_kernel_format( - backend=self.wna16_backend, - layer=layer, - quant_config=self.weight_quant, - input_dtype=self.marlin_input_dtype, - w13=layer.w13_weight_packed, - w2=layer.w2_weight_packed, - w13_scale=layer.w13_weight_scale, - w2_scale=layer.w2_weight_scale, - w13_g_idx=layer.w13_weight_g_idx, - w2_g_idx=layer.w2_weight_g_idx, - w13_qzeros=getattr(layer, "w13_weight_zero_point", None), - w2_qzeros=getattr(layer, "w2_weight_zero_point", None), - ) + ) = converted # Replace common parameters replace_parameter(layer, "w13_weight_packed", w13_qweight) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py index 1301c98f45b..23e3510614b 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py @@ -137,8 +137,16 @@ class CompressedTensorsW8A16Fp8(CompressedTensorsScheme): "weight_scale", convert_to_channelwise(layer.weight_scale, layer.logical_widths), ) + self.strategy = QuantizationStrategy.CHANNEL + self.weight_quant_key = STRATEGY_TO_WEIGHT_QUANT_KEY[self.strategy] + self.linear_kernel.config.weight_quant_key = self.weight_quant_key + # Canonicalize to (K, N) for the kernel. replace_parameter(layer, "weight", layer.weight.t()) + # Preserve the dim tags dropped by the transpose so layout-aware + # kernels see (K, N). + layer.weight.input_dim = 0 + layer.weight.output_dim = 1 self.linear_kernel.process_weights_after_loading(layer) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py index 1a240f6540d..e50778f0acb 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py @@ -181,6 +181,10 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme): # required by torch.compile to be torch.nn.Parameter layer.weight = Parameter(weight.data, requires_grad=False) + # Preserve the dim tags dropped by the transpose so layout-aware + # kernels (humming) see (K, N) instead of assuming (N, K). + layer.weight.input_dim = 0 + layer.weight.output_dim = 1 layer.weight_scale = Parameter(weight_scale.data, requires_grad=False) if input_scale is not None: layer.input_scale = Parameter(input_scale.data, requires_grad=False) diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 86818ed4b7e..626fc83cdff 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -714,6 +714,7 @@ class Fp8MoEMethod(FusedMoEMethodBase): fp8_backend=self.fp8_backend, experts_cls=self.experts_cls, routing_tables=layer._expert_routing_tables(), + layer=layer, ) def process_weights_after_loading(self, layer: RoutedExperts) -> None: @@ -786,6 +787,7 @@ class Fp8MoEMethod(FusedMoEMethodBase): swiglu_limit=getattr(layer, "swiglu_limit", None), gemm1_alpha=getattr(layer, "swiglu_alpha", None), gemm1_beta=getattr(layer, "swiglu_beta", None), + layer=layer, ) # Inject biases into the quant config if the model has them diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_ark_ops.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_ark_ops.py new file mode 100644 index 00000000000..cf54d1b531e --- /dev/null +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_ark_ops.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from functools import lru_cache +from typing import Any + +import torch + +from vllm.logger import init_logger +from vllm.platforms import current_platform +from vllm.utils.torch_utils import direct_register_custom_op + +logger = init_logger(__name__) + +_OPS_REGISTERED = False + + +@lru_cache(maxsize=1) +def get_ark_state() -> tuple[bool, str | None, Any | None, Any | None]: + """Return ARK availability, error details, cached module, and QuantLinear.""" + try: + import auto_round_kernel as ark + from auto_round_kernel.qlinear import QuantLinear + + logger.info("Successfully imported auto_round_kernel.") + except ImportError as error: + return False, str(error), None, None + + if getattr(ark, "cpu_lib", None) is None and getattr(ark, "xpu_lib", None) is None: + return ( + False, + "No ARK backend library is available.", + None, + None, + ) + + logger.info("Successfully loaded auto_round_kernel backend library.") + return True, None, ark, QuantLinear + + +def _inc_ark_woq_linear_impl( + x: torch.Tensor, + qweight: torch.Tensor, + bias: torch.Tensor | None, + out_features: int, + in_features: int, + group_size: int, + compute_type: str, + weight_type: str, + scale_type: str, + asym: bool, +) -> torch.Tensor: + ark = get_ark_state()[2] + assert ark is not None + + return ark.woqgemm_linear( + x, + qweight, + bias, + out_features, + in_features, + group_size, + compute_type, + weight_type, + scale_type, + asym, + ) + + +def _inc_ark_woq_linear_fake( + x: torch.Tensor, + qweight: torch.Tensor, + bias: torch.Tensor | None, + out_features: int, + in_features: int, + group_size: int, + compute_type: str, + weight_type: str, + scale_type: str, + asym: bool, +) -> torch.Tensor: + del qweight + del bias + del in_features + del group_size + del compute_type + del weight_type + del scale_type + del asym + return torch.empty( + (*x.shape[:-1], out_features), + dtype=x.dtype, + device=x.device, + ) + + +class ark_ops: + @staticmethod + def register_ops_once() -> None: + global _OPS_REGISTERED + if _OPS_REGISTERED: + return + + is_available, error_str, _, _ = get_ark_state() + if not is_available: + logger.debug( + "Skip registering ark op because ARK is unavailable: %s", + error_str or "unknown error", + ) + return + + direct_register_custom_op( + op_name="inc_ark_woq_linear", + op_func=_inc_ark_woq_linear_impl, + fake_impl=_inc_ark_woq_linear_fake, + dispatch_key=current_platform.dispatch_key, + ) + _OPS_REGISTERED = True + + +ark_ops.register_ops_once() + +__all__ = ["get_ark_state"] diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py index a212e4d3050..dd6c2fa2eaa 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py @@ -1,13 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from functools import lru_cache from typing import TYPE_CHECKING, Any import torch from torch.nn.parameter import Parameter -from vllm.logger import init_logger from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig from vllm.model_executor.layers.quantization.utils.marlin_utils import ( @@ -22,35 +20,10 @@ from vllm.scalar_type import scalar_types from .inc_scheme import INCLinearScheme -logger = init_logger(__name__) - if TYPE_CHECKING: from ..config_parser import INCLayerConfig -@lru_cache(maxsize=1) -def get_ark_state() -> tuple[bool, str | None, Any | None, Any | None]: - """Return ARK availability, error details, cached module, and QuantLinear.""" - try: - import auto_round_kernel as ark - from auto_round_kernel.qlinear import QuantLinear - - logger.info("Successfully imported auto_round_kernel.") - except ImportError as error: - return False, str(error), None, None - - if getattr(ark, "cpu_lib", None) is None and getattr(ark, "xpu_lib", None) is None: - return ( - False, - "No ARK backend library is available.", - None, - None, - ) - logger.info("Successfully loaded auto_round_kernel backend library.") - - return True, None, ark, QuantLinear - - class INCWNA16LinearScheme(INCLinearScheme): def __init__(self, layer_config: "INCLayerConfig") -> None: self.layer_config = layer_config @@ -380,6 +353,8 @@ class INCARKLinearMethod(INCXPULinearBase): def __init__(self, layer_config: "INCLayerConfig") -> None: super().__init__(layer_config) + from .inc_ark_ops import get_ark_state + is_available, error_str, _, quant_linear_cls = get_ark_state() if not is_available or quant_linear_cls is None: reason = error_str or "unknown error" @@ -453,9 +428,13 @@ class INCARKLinearMethod(INCXPULinearBase): ark_linear.bias.copy_(layer.bias.detach()) ark_linear.post_init() - layer.ark_linear = ark_linear - del layer.qweight + layer.qweight = Parameter(ark_linear.qweight.detach(), requires_grad=False) + layer.ark_bias = ark_linear.bias + layer.ark_compute_type = ark_linear.cdt + layer.ark_weight_type = ark_linear.wdt + layer.ark_scale_type = ark_linear.sdt + if hasattr(layer, "qzeros"): del layer.qzeros del layer.scales @@ -466,8 +445,18 @@ class INCARKLinearMethod(INCXPULinearBase): x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - del bias - return layer.ark_linear.forward(x) + return torch.ops.vllm.inc_ark_woq_linear.default( + x, + layer.qweight, + layer.ark_bias, + layer.out_features, + layer.in_features, + self.group_size, + layer.ark_compute_type, + layer.ark_weight_type, + layer.ark_scale_type, + not self.sym, + ) class INCXPUW4A16LinearScheme(INCXPULinearMethod): diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_scheme.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_scheme.py index e994b034944..80310358619 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_scheme.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_scheme.py @@ -36,10 +36,10 @@ class INCWna16Scheme(INCScheme): del config, layer if current_platform.is_xpu(): if layer_config.bits == 4 and layer_config.sym: + from .inc_ark_ops import get_ark_state from .inc_wna16_linear import ( INCARKLinearMethod, INCXPULinearMethod, - get_ark_state, ) is_ark_available, ark_error, _, _ = get_ark_state() @@ -57,10 +57,10 @@ class INCWna16Scheme(INCScheme): if current_platform.is_cpu() and layer_config.is_gptq: if layer_config.bits == 4 and layer_config.sym: + from .inc_ark_ops import get_ark_state from .inc_wna16_linear import ( INCARKLinearMethod, INCWNA16LinearScheme, - get_ark_state, ) is_ark_available, ark_error, _, _ = get_ark_state() diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 8fa1cb4d544..4a078f1357c 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -905,6 +905,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase): fp8_backend=self.fp8_backend, experts_cls=self.experts_cls, routing_tables=layer._expert_routing_tables(), + layer=layer, ) def process_weights_after_loading(self, layer: RoutedExperts) -> None: @@ -951,6 +952,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase): a1_scale=a1_scale, a2_scale=a2_scale, swiglu_limit=getattr(layer, "swiglu_limit", None), + layer=layer, ) def apply_monolithic( @@ -1602,7 +1604,9 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase): moe_quant_config=self.moe_quant_config, moe_config=self.moe, experts_cls=self.experts_cls, + backend=self.nvfp4_backend, routing_tables=layer._expert_routing_tables(), + layer=layer, ) self.moe_kernel.fused_experts.process_weights_after_loading(layer) @@ -1616,6 +1620,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase): a13_scale=layer.w13_input_scale, a2_scale=layer.w2_input_scale, swiglu_limit=getattr(layer, "swiglu_limit", None), + layer=layer, ) @property @@ -2162,6 +2167,7 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase): fp8_backend=self.mxfp8_backend, experts_cls=self.experts_cls, routing_tables=layer._expert_routing_tables(), + layer=layer, ) # No native MXFP8 MoE kernel on this device (e.g. gfx942): the emulation @@ -2207,6 +2213,7 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase): swiglu_limit=getattr(layer, "swiglu_limit", None), gemm1_alpha=getattr(layer, "swiglu_alpha", None), gemm1_beta=getattr(layer, "swiglu_beta", None), + layer=layer, ) def apply_monolithic( diff --git a/vllm/model_executor/layers/quantization/online/fp8.py b/vllm/model_executor/layers/quantization/online/fp8.py index 4d3a3158791..e63c285aa30 100644 --- a/vllm/model_executor/layers/quantization/online/fp8.py +++ b/vllm/model_executor/layers/quantization/online/fp8.py @@ -458,6 +458,7 @@ class _Fp8OnlineMoEBase(OnlineMoEMethodBase): fp8_backend=self.fp8_backend, experts_cls=self.experts_cls, routing_tables=layer._expert_routing_tables(), + layer=layer, ) def get_fused_moe_quant_config( @@ -486,6 +487,7 @@ class _Fp8OnlineMoEBase(OnlineMoEMethodBase): swiglu_limit=getattr(layer, "swiglu_limit", None), gemm1_alpha=getattr(layer, "swiglu_alpha", None), gemm1_beta=getattr(layer, "swiglu_beta", None), + layer=layer, ) diff --git a/vllm/model_executor/layers/quantization/online/int8.py b/vllm/model_executor/layers/quantization/online/int8.py index 4274b16d55a..ef76594164a 100644 --- a/vllm/model_executor/layers/quantization/online/int8.py +++ b/vllm/model_executor/layers/quantization/online/int8.py @@ -13,6 +13,7 @@ if TYPE_CHECKING: from vllm.model_executor.layers.fused_moe import RoutedExperts from vllm.model_executor.layers.fused_moe.oracle.int8 import ( + convert_to_int8_moe_kernel_format, make_int8_moe_kernel, make_int8_moe_quant_config, select_int8_moe_backend, @@ -92,22 +93,36 @@ class Int8OnlineMoEMethod(OnlineMoEMethodBase): replace_parameter(layer, "w2_scale", w2_scale) def _setup_kernel(self, layer: RoutedExperts) -> None: + w13, w2 = convert_to_int8_moe_kernel_format( + int8_backend=self.int8_backend, + w13=layer.w13_weight, + w2=layer.w2_weight, + layer=layer, + w13_scale=layer.w13_scale, + ) + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w2_weight", w2) + self.moe_quant_config = self.get_fused_moe_quant_config(layer) assert self.moe_quant_config is not None assert self.experts_cls is not None self.moe_kernel = make_int8_moe_kernel( + int8_backend=self.int8_backend, moe_quant_config=self.moe_quant_config, moe_config=self.moe, experts_cls=self.experts_cls, routing_tables=layer._expert_routing_tables(), + layer=layer, ) def get_fused_moe_quant_config( self, layer: torch.nn.Module ) -> "FusedMoEQuantConfig | None": return make_int8_moe_quant_config( - w1_scale=layer.w13_scale, - w2_scale=layer.w2_scale, + int8_backend=self.int8_backend, + w1_scale=getattr(layer, "w13_scale", None), + w2_scale=getattr(layer, "w2_scale", None), w1_bias=getattr(layer, "w13_bias", None), w2_bias=getattr(layer, "w2_bias", None), + layer=layer, ) diff --git a/vllm/model_executor/layers/quantization/online/mxfp8.py b/vllm/model_executor/layers/quantization/online/mxfp8.py index 09d581a0734..84a81bd9064 100644 --- a/vllm/model_executor/layers/quantization/online/mxfp8.py +++ b/vllm/model_executor/layers/quantization/online/mxfp8.py @@ -200,6 +200,7 @@ class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase): fp8_backend=self.fp8_backend, experts_cls=self.experts_cls, routing_tables=layer._expert_routing_tables(), + layer=layer, ) def get_fused_moe_quant_config( @@ -226,6 +227,7 @@ class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase): swiglu_limit=getattr(layer, "swiglu_limit", None), gemm1_alpha=getattr(layer, "swiglu_alpha", None), gemm1_beta=getattr(layer, "swiglu_beta", None), + layer=layer, ) def process_weights_after_loading(self, layer: Module) -> None: diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index bce888415ce..7bdf963b512 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -1562,7 +1562,9 @@ class QuarkNvfp4MoEMethod(QuarkMoEMethod): moe_quant_config=self.moe_quant_config, moe_config=self.moe, experts_cls=self.experts_cls, + backend=self.nvfp4_backend, routing_tables=layer._expert_routing_tables(), + layer=layer, ) def get_fused_moe_quant_config( @@ -1576,6 +1578,7 @@ class QuarkNvfp4MoEMethod(QuarkMoEMethod): w2_scale_2=layer.w2_weight_scale_2, a13_scale=layer.w13_input_scale_2, a2_scale=layer.w2_input_scale_2, + layer=layer, ) def apply( diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index 5b77ac39225..632cca1fda2 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from enum import Enum from typing import TYPE_CHECKING import torch @@ -15,12 +14,6 @@ if TYPE_CHECKING: logger = init_logger(__name__) -class FlashinferMoeBackend(Enum): - TENSORRT_LLM = "TensorRT-LLM" - CUTLASS = "CUTLASS" - CUTEDSL = "CUTEDSL" - - def activation_to_flashinfer_int(activation: MoEActivation) -> int: return activation_to_flashinfer_type(activation).value @@ -97,16 +90,6 @@ def rotate_weights_for_fi_trtllm_fp8_per_tensor_moe( ) -def is_flashinfer_supporting_global_sf(backend: FlashinferMoeBackend | None) -> bool: - # TODO(shuw@nvidia): Update when new backends are added. - backends_supporting_global_sf = ( - FlashinferMoeBackend.CUTLASS, - FlashinferMoeBackend.TENSORRT_LLM, - FlashinferMoeBackend.CUTEDSL, - ) - return backend in backends_supporting_global_sf - - def convert_moe_weights_to_flashinfer_trtllm_block_layout( cache_permute_indices: dict[torch.Size, torch.Tensor], w13_weight: torch.Tensor, diff --git a/vllm/model_executor/layers/quantization/utils/humming_utils.py b/vllm/model_executor/layers/quantization/utils/humming_utils.py index d84a2e12f54..2d9e7aca766 100644 --- a/vllm/model_executor/layers/quantization/utils/humming_utils.py +++ b/vllm/model_executor/layers/quantization/utils/humming_utils.py @@ -764,7 +764,7 @@ def _convert_sublayer_to_humming( Returns: Tuple of (converted_weight_schema, converted_input_schema) """ - from humming.schema import HummingWeightSchema + from vllm.utils.humming import HummingWeightSchema if isinstance(weight_schema, HummingWeightSchema): # Already in Humming format @@ -814,7 +814,7 @@ def _prepare_and_transform_sublayer( This calls Humming's prepare_layer_meta and transform_humming_layer. """ - from humming.layer import HummingMethod + from vllm.utils.humming import HummingMethod HummingMethod.prepare_layer_meta( layer=layer, @@ -866,7 +866,7 @@ def _process_single_sublayer( Returns: Tuple of (final_weight_schema, final_input_schema) """ - from humming.schema import HummingWeightSchema + from vllm.utils.humming import HummingWeightSchema # Step 1: Convert from checkpoint format to humming format if needed current_weight_schema, current_input_schema = _convert_sublayer_to_humming( @@ -958,12 +958,10 @@ def convert_to_humming_moe_kernel_format( "Must provide either weight_schema/input_schema or quant_config" ) - from humming.layer import HummingInputSchema - from humming.schema import BaseWeightSchema - from vllm.model_executor.layers.quantization.utils.humming_utils import ( humming_is_layer_skipped, ) + from vllm.utils.humming import BaseWeightSchema, HummingInputSchema if weight_schema is None: weight_schema = BaseWeightSchema.from_config(quant_config) diff --git a/vllm/model_executor/layers/quantization/utils/int8_utils.py b/vllm/model_executor/layers/quantization/utils/int8_utils.py index eac6b11b219..e0db2526948 100644 --- a/vllm/model_executor/layers/quantization/utils/int8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/int8_utils.py @@ -16,43 +16,6 @@ from vllm.triton_utils import tl, triton logger = logging.getLogger(__name__) -def apply_w8a8_block_int8_linear( - input: torch.Tensor, - weight: torch.Tensor, - block_size: list[int], - weight_scale: torch.Tensor, - input_scale: torch.Tensor | None = None, - bias: torch.Tensor | None = None, -) -> torch.Tensor: - assert input_scale is None - # View input as 2D matrix for fp8 methods - input_2d = input.view(-1, input.shape[-1]) - output_shape = [*input.shape[:-1], weight.shape[0]] - - q_input, x_scale = per_token_group_quant_int8(input_2d, block_size[1]) - output = w8a8_block_int8_matmul( - q_input, weight, x_scale, weight_scale, block_size, output_dtype=input.dtype - ) - - if bias is not None: - output = output + bias - return output.to(dtype=input.dtype).view(*output_shape) - - -def input_to_int8( - x: torch.Tensor, dtype: torch.dtype = torch.int8 -) -> tuple[torch.Tensor, torch.Tensor]: - """This function quantizes input values to int8 values with - tensor-wise quantization.""" - iinfo = torch.iinfo(dtype) - min_val, max_val = x.aminmax() - amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12) - int8_min, int8_max = iinfo.min, iinfo.max - scale = int8_max / amax - x_scl_sat = (x * scale).clamp(min=int8_min, max=int8_max) - return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal() - - def block_dequant( x_q_block: torch.Tensor, x_s: torch.Tensor, diff --git a/vllm/model_executor/layers/quantization/utils/machete_utils.py b/vllm/model_executor/layers/quantization/utils/machete_utils.py index 95d8102ea50..7bf71ca99c5 100644 --- a/vllm/model_executor/layers/quantization/utils/machete_utils.py +++ b/vllm/model_executor/layers/quantization/utils/machete_utils.py @@ -16,10 +16,6 @@ def query_machete_supported_quant_types(zero_points: bool) -> list[ScalarType]: return [scalar_types.uint4b8, scalar_types.uint8b128] -def query_machete_supported_act_types(zero_points: bool) -> list[ScalarType]: - return [torch.float16, torch.bfloat16] - - def query_machete_supported_group_sizes(act_type: torch.dtype) -> list[int]: """ Queries the supported group sizes for Machete based on the activation type. diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils.py b/vllm/model_executor/layers/quantization/utils/marlin_utils.py index ea47ed06cbf..f6d96f574d8 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils.py @@ -738,71 +738,3 @@ def apply_gptq_marlin_linear( output = marlin_unpad_output(output, output_size_per_partition, padded_n) return output.reshape(out_shape) - - -def apply_awq_marlin_linear( - input: torch.Tensor, - weight: torch.Tensor, - weight_scale: torch.Tensor, - weight_zp: torch.Tensor, - g_idx: torch.Tensor, - g_idx_sort_indices: torch.Tensor, - workspace: torch.Tensor, - quant_type: ScalarType, - output_size_per_partition: int, - input_size_per_partition: int, - input_global_scale: torch.Tensor | None = None, - bias: torch.Tensor | None = None, - use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT, - input_dtype: torch.dtype | None = None, -) -> torch.Tensor: - reshaped_x = input.reshape(-1, input.shape[-1]) - out_shape = input.shape[:-1] + (output_size_per_partition,) - - padded_n, padded_k = marlin_repacked_nk(weight, quant_type.size_bits) - reshaped_x = marlin_pad_dim(reshaped_x, input_size_per_partition, padded_k) - - use_atomic_add = should_use_atomic_add_reduce( - m=reshaped_x.size(0), - n=padded_n, - k=padded_k, - device=input.device, - dtype=input.dtype, - ) - - a_scales = None - if input_dtype == torch.int8: - assert quant_type == scalar_types.uint4, ( - "W8A8-INT8 is not supported by marlin kernel." - ) - reshaped_x, a_scales = marlin_quant_input(reshaped_x, input_dtype) - a_scales = a_scales * input_global_scale - elif input_dtype == torch.float8_e4m3fn: - assert quant_type == scalar_types.uint4, ( - "INT8 weight + FP8 activation is not supported." - ) - reshaped_x, a_scales = marlin_quant_input(reshaped_x, input_dtype) - - output = ops.marlin_gemm( - reshaped_x, - None, - weight, - bias, - weight_scale, - a_scales, - None, - weight_zp, - g_idx, - g_idx_sort_indices, - workspace, - quant_type, - size_m=reshaped_x.shape[0], - size_n=padded_n, - size_k=padded_k, - use_atomic_add=use_atomic_add, - use_fp32_reduce=use_fp32_reduce, - is_zp_float=False, - ) - - output = marlin_unpad_output(output, output_size_per_partition, padded_n) - return output.reshape(out_shape) diff --git a/vllm/model_executor/layers/quantization/utils/quant_utils.py b/vllm/model_executor/layers/quantization/utils/quant_utils.py index 0c5cbae2a4f..705da43c373 100644 --- a/vllm/model_executor/layers/quantization/utils/quant_utils.py +++ b/vllm/model_executor/layers/quantization/utils/quant_utils.py @@ -216,6 +216,12 @@ kInt4W4A8StaticGroupSym = QuantKey( torch.int8, kInt4W4A8StaticGroupScale, symmetric=True ) +kInt4W4A8StaticChannelSym = QuantKey( + torch.int8, + ScaleDesc(torch.float32, True, GroupShape.PER_CHANNEL), + symmetric=True, +) + def create_fp8_quant_key( static: bool, diff --git a/vllm/model_executor/models/bailing_moe_linear.py b/vllm/model_executor/models/bailing_moe_linear.py index 7bf5fb04077..5cc057c3b45 100644 --- a/vllm/model_executor/models/bailing_moe_linear.py +++ b/vllm/model_executor/models/bailing_moe_linear.py @@ -533,6 +533,10 @@ class BailingMoeV25Model(nn.Module): def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.word_embeddings(input_ids) + @property + def embed_tokens(self) -> nn.Module: + return self.word_embeddings + def forward( self, input_ids: torch.Tensor | None, diff --git a/vllm/model_executor/models/bailing_moe_mtp.py b/vllm/model_executor/models/bailing_moe_mtp.py new file mode 100644 index 00000000000..da6b1ddb8b6 --- /dev/null +++ b/vllm/model_executor/models/bailing_moe_mtp.py @@ -0,0 +1,380 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only Bailing MoE v2.5 MTP model.""" + +from collections.abc import Iterable + +import torch +import torch.nn as nn +from transformers.configuration_utils import PretrainedConfig + +from vllm.compilation.decorators import support_torch_compile +from vllm.config import VllmConfig +from vllm.distributed import get_pp_group +from vllm.model_executor.layers.fused_moe import ( + fused_moe_make_expert_params_mapping, +) +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.model_executor.models.bailing_moe_linear import ( + BailingMoeV25, + BailingMoeV25MLAAttention, +) +from vllm.sequence import IntermediateTensors + +from .utils import PPMissingLayer, is_pp_missing_parameter, maybe_prefix + + +def _get_draft_hf_config(vllm_config: VllmConfig) -> PretrainedConfig: + speculative_config = vllm_config.speculative_config + if speculative_config is not None: + draft_model_config = speculative_config.draft_model_config + if draft_model_config is not None: + return draft_model_config.hf_config + return vllm_config.model_config.hf_config + + +class BailingMTPSharedHead(nn.Module): + def __init__( + self, + config: PretrainedConfig, + prefix: str, + vllm_config: VllmConfig, + ) -> None: + super().__init__() + self.head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=vllm_config.quant_config, + prefix=maybe_prefix(prefix, "head"), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return hidden_states + + +class BailingMoeV25MultiTokenPredictorLayer(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + prefix: str, + layer_id: int, + ) -> None: + super().__init__() + config = _get_draft_hf_config(vllm_config) + self.config = config + self.layer_id = layer_id + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.self_attn = BailingMoeV25MLAAttention( + config, + quant_config=vllm_config.quant_config, + layer_id=layer_id, + prefix=maybe_prefix(prefix, "self_attn"), + cache_config=vllm_config.cache_config, + ) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.mlp = BailingMoeV25( + config, + quant_config=vllm_config.quant_config, + layer_id=layer_id, + prefix=maybe_prefix(prefix, "mlp"), + ) + self.final_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.shared_head = BailingMTPSharedHead( + config, + maybe_prefix(prefix, "shared_head"), + vllm_config, + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_index: int = 0, + ) -> torch.Tensor: + assert inputs_embeds is not None + inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds) + inputs_embeds = self.enorm(inputs_embeds) + previous_hidden_states = self.hnorm(previous_hidden_states) + + hidden_states = self.eh_proj( + torch.cat([inputs_embeds, previous_hidden_states], dim=-1) + ) + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + hidden_states = self.self_attn(hidden_states, positions) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states.to(residual.device) + return self.final_layernorm(hidden_states) + + +class BailingMoeV25MultiTokenPredictor(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + config = _get_draft_hf_config(vllm_config) + self.mtp_start_layer_idx = config.num_hidden_layers + self.num_mtp_layers = config.num_nextn_predict_layers + self.layers = nn.ModuleDict( + { + str(idx): BailingMoeV25MultiTokenPredictorLayer( + vllm_config, + f"{prefix}.layers.{idx}", + idx, + ) + for idx in range( + self.mtp_start_layer_idx, + self.mtp_start_layer_idx + self.num_mtp_layers, + ) + } + ) + if get_pp_group().is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size, + prefix=maybe_prefix(prefix, "embed_tokens"), + ) + else: + self.embed_tokens = PPMissingLayer() + self.logits_processor = LogitsProcessor(config.vocab_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, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + current_step_idx = spec_step_idx % self.num_mtp_layers + return self.layers[str(self.mtp_start_layer_idx + current_step_idx)]( + input_ids, + positions, + previous_hidden_states, + inputs_embeds, + current_step_idx, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + lm_head: nn.Module | None = None, + ) -> torch.Tensor: + current_step_idx = spec_step_idx % self.num_mtp_layers + mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)] + head = lm_head if lm_head is not None else mtp_layer.shared_head.head + return self.logits_processor( + head, + mtp_layer.shared_head(hidden_states), + ) + + +@support_torch_compile +class BailingMoeV25MTPModel(nn.Module): + packed_modules_mapping = { + "gate_up_proj": ["gate_proj", "up_proj"], + "fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"], + } + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + self.config = _get_draft_hf_config(vllm_config) + self.lm_head: nn.Module | None = None + self.model = BailingMoeV25MultiTokenPredictor( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "model"), + ) + + 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, + hidden_states: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + return self.model( + input_ids, + positions, + hidden_states, + inputs_embeds, + spec_step_idx, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + return self.model.compute_logits(hidden_states, spec_step_idx, self.lm_head) + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + return fused_moe_make_expert_params_mapping( + self, + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.num_experts, + num_redundant_experts=0, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + (".fused_qkv_a_proj", ".q_a_proj", 0), + (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + expert_params_mapping = list(self.get_expert_mapping()) + params_dict = dict(self.named_parameters(remove_duplicate=False)) + loaded_params: set[str] = set() + loaded_mtp_layers: set[int] = set() + + def load_param( + name: str, + loaded_weight: torch.Tensor, + shard_id=None, + ) -> bool: + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + return False + if name not in params_dict or is_pp_missing_parameter(name, self): + return False + + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + if shard_id is None: + weight_loader(param, loaded_weight) + elif isinstance(shard_id, int): + weight_loader(param, loaded_weight, shard_id) + else: + weight_loader( + param, + loaded_weight, + name, + expert_id=shard_id[0], + shard_id=shard_id[1], + ) + loaded_params.add(name) + return True + + def get_spec_layer_idx(name: str) -> int | None: + if not name.startswith("model.layers."): + return None + try: + layer_idx = int(name.split("model.layers.", 1)[1].split(".", 1)[0]) + except (IndexError, ValueError): + return None + mtp_idx = layer_idx - self.config.num_hidden_layers + if 0 <= mtp_idx < self.config.num_nextn_predict_layers: + return layer_idx + return None + + def normalize_name(name: str) -> str: + name = name.replace(".attention.dense", ".self_attn.o_proj") + name = name.replace(".attention.", ".self_attn.") + return name.replace( + "mlp.gate.e_score_correction_bias", + "mlp.gate.expert_bias", + ) + + def load_lm_head(loaded_weight: torch.Tensor) -> None: + for layer_idx in range( + self.model.mtp_start_layer_idx, + self.model.mtp_start_layer_idx + self.model.num_mtp_layers, + ): + name = f"model.layers.{layer_idx}.shared_head.head.weight" + load_param(name, loaded_weight) + + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + + if name == "model.word_embeddings.weight": + load_param("model.embed_tokens.weight", loaded_weight) + continue + if name == "lm_head.weight": + load_lm_head(loaded_weight) + continue + + spec_layer = get_spec_layer_idx(name) + if spec_layer is None: + continue + name = normalize_name(name) + + loaded = False + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + if "mlp.experts." in name and name not in params_dict: + continue + mapped_name = name.replace(weight_name, param_name) + if load_param(mapped_name, loaded_weight, shard_id): + loaded = True + break + if loaded: + loaded_mtp_layers.add(spec_layer) + continue + + if "mlp.experts" in name: + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in name: + continue + mapped_name = name.replace(weight_name, param_name) + if load_param( + mapped_name, + loaded_weight, + (expert_id, shard_id), + ): + loaded = True + break + if loaded: + loaded_mtp_layers.add(spec_layer) + continue + + if load_param(name, loaded_weight): + loaded_mtp_layers.add(spec_layer) + + for layer_idx in range( + self.model.mtp_start_layer_idx, + self.model.mtp_start_layer_idx + self.model.num_mtp_layers, + ): + if layer_idx not in loaded_mtp_layers: + raise ValueError( + f"Bailing MTP speculative decoding layer {layer_idx} " + "weights are missing from checkpoint. Use a checkpoint " + "that includes MTP layer weights, or disable speculative " + "decoding." + ) + return loaded_params diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 539e912327d..134a6d84ed4 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -829,7 +829,7 @@ def _try_load_fp8_indexer_wk( if "indexer.wk." not in name or "wk_weights" in name: return False # Weight is not an isolated WK weight for the indexer, ignore. is_weight = name.endswith(".weight") and tensor.dtype == torch.float8_e4m3fn - is_scale = "weight_scale_inv" in name + is_scale = "weight_scale" in name if not is_weight and not is_scale: return False # WK is not in FP8 format, ignore. # Buffer this tensor (weight or scale) until both have arrived. @@ -1511,7 +1511,10 @@ class DeepseekV2Model(nn.Module): ("qkv_proj", "v_proj", "v"), ] # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj) - _pending_wk_fp8: dict = {} # When WK is in FP8, we dequant to BF16 for fusion + _pending_wk_fp8 = getattr(self, "_pending_indexer_wk_fp8", None) + if _pending_wk_fp8 is None: + self._pending_indexer_wk_fp8 = _pending_wk_fp8 = {} + indexer_fused_mapping = [ ("wk_weights_proj", "wk", 0), ("wk_weights_proj", "weights_proj", 1), diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index eebb5ef148e..11a10131df1 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -47,6 +47,7 @@ from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.model_executor.models.transformers.utils import recursive_replace_linear from vllm.model_executor.models.utils import WeightsMapper, maybe_prefix from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.platforms import current_platform from vllm.v1.outputs import LogprobsTensors from vllm.v1.worker.gpu.attn_utils import build_attn_metadata from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor, async_copy_to_gpu @@ -516,11 +517,6 @@ def _compiled_sample_step( num_decode = decode_slots.shape[0] device = decode_slots.device - # Clear outputs so prefill / non-decode slots report 0 (decode slots are - # overwritten below). - sampled.zero_() - num_sampled.zero_() - # ---- Phase 1: Temperature schedule ---- steps_f = step_tensor[decode_slots].float() remaining = (max_denoising_steps - steps_f).clamp(min=1.0) @@ -800,7 +796,10 @@ class DiffusionGemmaModelState(ModelState): max_denoising_steps=max_denoising_steps, device=device, hidden_size=text_config.hidden_size, - stability_threshold=self.gen_config["stability_threshold"], + # In Transformers, `stability_threshold=1` (the default) means the current + # step must match the previous step. In vLLM, the history buffer includes + # the current step, so we add 1 to match the same behavior. + stability_threshold=self.gen_config["stability_threshold"] + 1, ) self._req_id_to_index: dict[str, int] = {} @@ -1273,9 +1272,12 @@ class DiffusionSampler: src = (starts.unsqueeze(1) + ar.unsqueeze(0)).clamp_max(logits.shape[0] - 1) logits = logits[src.reshape(-1)] * valid.reshape(-1, 1).to(logits.dtype) - # Cleared inside _compiled_sample_step so prefill/non-decode slots stay 0. + # Clear once: the tiled loop below only scatters its own decode slots, + # so it must not re-clear earlier tiles' writes. sampled = self._sampled[:num_reqs] num_sampled = self._num_sampled[:num_reqs] + sampled.zero_() + num_sampled.zero_() all_slots = input_batch.idx_mapping[:num_reqs] @@ -1283,94 +1285,109 @@ class DiffusionSampler: # since it mutates is_encoder_phase (commit→False, converge→True). is_committing = states.is_encoder_phase[decode_slots].clone() - # --- Single compiled call: temp → sample → probs → post-process --- - scaled = _compiled_sample_step( - logits, - decode_slots, - decode_idx, - all_slots, - valid_canvas_len, - # State - states.canvas, - states.argmax_canvas, - states.step, - states.is_encoder_phase, - states.confident, - states.self_conditioning_embeds, - self.embed_weight, - self.normalizer, - states.accepted_canvas_history, - states.accepted_canvas_history_len, - # Output - sampled, - num_sampled, - self.req_states.draft_tokens, - # Config - max_denoising_steps=float(states.max_denoising_steps), - t_min=self.t_min, - t_max=self.t_max, - confidence_threshold=self.confidence_threshold, - vocab_size=self.vocab_size, - CL=self.canvas_length, - ST=states.stability_threshold, - entropy_bound=self.entropy_bound, - sc_vocab_start=self.sc_vocab_start, - sc_vocab_end=self.sc_vocab_end, - tp_size=self.tp_size, - tp_group_name=self.tp_group_name, - ) - - # --- Logprobs: stash on convergence, return on commit --- slots_np = input_batch.idx_mapping_np[:num_reqs] is_decode_np = per_req_nlogits_np > 0 - - logprobs_tensors = None max_num_logprobs = self.sampling_states.max_num_logprobs(slots_np) - if max_num_logprobs >= 0: - # Denoise steps that just converged: the compiled step flipped - # is_encoder_phase from False→True. Detect as slots where - # is_encoder_phase is now True but is_committing was False. - converged_mask = states.is_encoder_phase[decode_slots] - just_converged = converged_mask & ~is_committing - if just_converged.any(): - flat_logits = scaled.reshape(-1, scaled.shape[-1]) - argmax_tokens = scaled.argmax(dim=-1) - for local_idx in just_converged.nonzero(as_tuple=True)[0]: - li = local_idx.item() - slot = decode_slots[local_idx] - # Stash only the real canvas positions (== CL unless this - # canvas was truncated near max_model_len); padded tail - # positions are never emitted. - k_i = int(valid_canvas_len_np[li]) - start = li * CL - self._pending_logprobs[slot.item()] = compute_topk_logprobs( - flat_logits[start : start + k_i], - max_num_logprobs, - argmax_tokens[local_idx][:k_i], - ) - # Commit steps: is_committing was True at entry. Reassemble - # previously stashed logprobs and attach to SamplerOutput. - if is_committing.any() and self._pending_logprobs: - parts_ids, parts_lp, parts_ranks = [], [], [] - cu_gen: list[int] = [] - flat_offset = 0 - for i in range(num_reqs): - cu_gen.append(flat_offset) - slot = int(slots_np[i]) - if is_decode_np[i] and slot in self._pending_logprobs: - lp = self._pending_logprobs.pop(slot) - parts_ids.append(lp.logprob_token_ids) - parts_lp.append(lp.logprobs) - parts_ranks.append(lp.selected_token_ranks) - flat_offset += lp.logprobs.shape[0] - if parts_ids: - logprobs_tensors = LogprobsTensors( - logprob_token_ids=torch.cat(parts_ids), - logprobs=torch.cat(parts_lp), - selected_token_ranks=torch.cat(parts_ranks), - cu_num_generated_tokens=cu_gen, - ) + # Sample over the [num_decode * CL, vocab] logits. The fp32 pipeline in + # _compiled_sample_step keeps several live [group * CL, vocab] copies, so + # size each tile to a fraction of free memory to bound the transient at + # high concurrency. Tiling is bit-identical to a single pass. + group = max(num_decode, 1) + if num_decode > 0: + free, _ = current_platform.mem_get_info() + # ~10 transient fp32 copies of [group * CL, vocab] inside the step + # (eager peaks at ~8; pad for allocator overhead and small tensors). + bytes_per_req = CL * self.vocab_size * 4 * 10 + budget = int(free * 0.5) // max(bytes_per_req, 1) + group = max(1, min(num_decode, budget)) + + for start_req in range(0, num_decode, group): + end_req = min(start_req + group, num_decode) + tile = slice(start_req, end_req) + tile_slots = decode_slots[tile] + + scaled = _compiled_sample_step( + logits[start_req * CL : end_req * CL], + tile_slots, + decode_idx[tile], + all_slots, + valid_canvas_len[tile], + # State + states.canvas, + states.argmax_canvas, + states.step, + states.is_encoder_phase, + states.confident, + states.self_conditioning_embeds, + self.embed_weight, + self.normalizer, + states.accepted_canvas_history, + states.accepted_canvas_history_len, + # Output + sampled, + num_sampled, + self.req_states.draft_tokens, + # Config + max_denoising_steps=float(states.max_denoising_steps), + t_min=self.t_min, + t_max=self.t_max, + confidence_threshold=self.confidence_threshold, + vocab_size=self.vocab_size, + CL=CL, + ST=states.stability_threshold, + entropy_bound=self.entropy_bound, + sc_vocab_start=self.sc_vocab_start, + sc_vocab_end=self.sc_vocab_end, + tp_size=self.tp_size, + tp_group_name=self.tp_group_name, + ) + + # Logprobs for denoise steps that just converged (is_encoder_phase + # flipped False→True), stashed per tile so `scaled` is freed each tile. + if max_num_logprobs >= 0: + converged_mask = states.is_encoder_phase[tile_slots] + just_converged = converged_mask & ~is_committing[tile] + if just_converged.any(): + flat_logits = scaled.reshape(-1, scaled.shape[-1]) + argmax_tokens = scaled.argmax(dim=-1) + for local_idx in just_converged.nonzero(as_tuple=True)[0]: + li = local_idx.item() + slot = tile_slots[local_idx] + # Stash only the real canvas positions (== CL unless this + # canvas was truncated near max_model_len); padded tail + # positions are never emitted. + k_i = int(valid_canvas_len_np[start_req + li]) + pos = li * CL + self._pending_logprobs[slot.item()] = compute_topk_logprobs( + flat_logits[pos : pos + k_i], + max_num_logprobs, + argmax_tokens[local_idx][:k_i], + ) + + # Commit steps: is_committing was True at entry. Reassemble previously + # stashed logprobs and attach to SamplerOutput. + logprobs_tensors = None + if max_num_logprobs >= 0 and is_committing.any() and self._pending_logprobs: + parts_ids, parts_lp, parts_ranks = [], [], [] + cu_gen: list[int] = [] + flat_offset = 0 + for i in range(num_reqs): + cu_gen.append(flat_offset) + slot = int(slots_np[i]) + if is_decode_np[i] and slot in self._pending_logprobs: + lp = self._pending_logprobs.pop(slot) + parts_ids.append(lp.logprob_token_ids) + parts_lp.append(lp.logprobs) + parts_ranks.append(lp.selected_token_ranks) + flat_offset += lp.logprobs.shape[0] + if parts_ids: + logprobs_tensors = LogprobsTensors( + logprob_token_ids=torch.cat(parts_ids), + logprobs=torch.cat(parts_lp), + selected_token_ranks=torch.cat(parts_ranks), + cu_num_generated_tokens=cu_gen, + ) return self._build_output( input_batch, diff --git a/vllm/model_executor/models/gemma4_mtp.py b/vllm/model_executor/models/gemma4_mtp.py index b30ab4b7aef..f4d200d11d9 100644 --- a/vllm/model_executor/models/gemma4_mtp.py +++ b/vllm/model_executor/models/gemma4_mtp.py @@ -51,6 +51,7 @@ from .utils import ( AutoWeightsLoader, WeightsMapper, extract_layer_index, + get_draft_quant_config, maybe_prefix, ) @@ -182,14 +183,14 @@ class Gemma4MTPAttention(nn.Module): hidden_size, self.total_num_heads * self.head_dim, bias=config.attention_bias, - quant_config=None, + quant_config=quant_config, prefix=f"{prefix}.q_proj", ) self.o_proj = RowParallelLinear( self.total_num_heads * self.head_dim, hidden_size, bias=config.attention_bias, - quant_config=None, + quant_config=quant_config, prefix=f"{prefix}.o_proj", ) self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) @@ -304,7 +305,7 @@ class Gemma4MTPDecoderLayer(nn.Module): hidden_size=self.hidden_size, intermediate_size=text_config.intermediate_size, hidden_activation=text_config.hidden_activation, - quant_config=None, + quant_config=quant_config, prefix=f"{prefix}.mlp", ) @@ -357,7 +358,9 @@ class Gemma4MultiTokenPredictor(nn.Module): config = vllm_config.speculative_config.draft_model_config.hf_config text_config = _get_text_config(config) + quant_config = get_draft_quant_config(vllm_config) self.config = text_config + self.quant_config = quant_config self.hidden_size = text_config.hidden_size self.backbone_hidden_size = getattr( @@ -369,6 +372,8 @@ class Gemma4MultiTokenPredictor(nn.Module): self.embed_tokens = VocabParallelEmbedding( self.vocab_size, self.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", ) self.pre_projection = ColumnParallelLinear( @@ -376,6 +381,7 @@ class Gemma4MultiTokenPredictor(nn.Module): self.hidden_size, bias=False, gather_output=True, + quant_config=quant_config, prefix=f"{prefix}.pre_projection", ) @@ -384,6 +390,7 @@ class Gemma4MultiTokenPredictor(nn.Module): self.backbone_hidden_size, bias=False, input_is_parallel=False, + quant_config=quant_config, prefix=f"{prefix}.post_projection", ) @@ -391,7 +398,7 @@ class Gemma4MultiTokenPredictor(nn.Module): Gemma4MTPDecoderLayer( text_config, cache_config=vllm_config.cache_config, - quant_config=vllm_config.quant_config, + quant_config=quant_config, prefix=f"{prefix}.layers.{idx}", ) for idx in range(self.num_mtp_layers) @@ -473,6 +480,7 @@ class Gemma4MTP(nn.Module): super().__init__() config = vllm_config.speculative_config.draft_model_config.hf_config text_config = _get_text_config(config) + self.quant_config = get_draft_quant_config(vllm_config) self.config = config self._stable_full_lm_head_weight: torch.Tensor | None = None @@ -488,6 +496,7 @@ class Gemma4MTP(nn.Module): self.lm_head = ParallelLMHead( text_config.vocab_size, text_config.hidden_size, + quant_config=self.quant_config, prefix=maybe_prefix(prefix, "lm_head"), ) if getattr(config, "tie_word_embeddings", True): diff --git a/vllm/model_executor/models/kimi_k25.py b/vllm/model_executor/models/kimi_k25.py index 7321b913605..e2b74744b41 100644 --- a/vllm/model_executor/models/kimi_k25.py +++ b/vllm/model_executor/models/kimi_k25.py @@ -56,6 +56,10 @@ from vllm.sequence import IntermediateTensors from vllm.transformers_utils.configs.kimi_k25 import KimiK25Config from vllm.transformers_utils.processor import cached_get_image_processor from vllm.transformers_utils.processors.kimi_k25 import KimiK25Processor +from vllm.transformers_utils.processors.kimi_k25_vision_fused import ( + KimiK25FusedVisionProcessor, +) +from vllm.utils.import_utils import is_numba_available from vllm.utils.tensor_schema import TensorSchema, TensorShape from .utils import ( @@ -108,10 +112,16 @@ class KimiK25ProcessingInfo(BaseProcessingInfo): self.hf_config = hf_config = self.get_hf_config() tokenizer = self.get_tokenizer() + processor_cls = KimiK25FusedVisionProcessor if is_numba_available() else None + logger.info_once( + "Using %s image preprocessing for Kimi-K2.5/K2.6 vision chunks.", + "fused CPU" if processor_cls is not None else "remote HF", + ) image_processor = cached_get_image_processor( self.ctx.model_config.model, revision=self.ctx.model_config.revision, trust_remote_code=self.ctx.model_config.trust_remote_code, + processor_cls_overrides=processor_cls, ) # Resolve token ID from the tokenizer because transformers v5 diff --git a/vllm/model_executor/models/llama_eagle3.py b/vllm/model_executor/models/llama_eagle3.py index 8ce86c77807..2d859bd4918 100644 --- a/vllm/model_executor/models/llama_eagle3.py +++ b/vllm/model_executor/models/llama_eagle3.py @@ -233,7 +233,10 @@ class LlamaModel(nn.Module): ) -> tuple[torch.Tensor, torch.Tensor]: if input_embeds is None: input_embeds = self.embed_input_ids(input_ids) - assert hidden_states.shape[-1] == input_embeds.shape[-1] + torch._assert( + hidden_states.shape[-1] == input_embeds.shape[-1], + "hidden_states and input_embeds must have the same last dimension", + ) residual = None for layer in self.layers: diff --git a/vllm/model_executor/models/moss_audio.py b/vllm/model_executor/models/moss_audio.py index fdfe982fca3..bef8057d2fb 100644 --- a/vllm/model_executor/models/moss_audio.py +++ b/vllm/model_executor/models/moss_audio.py @@ -862,6 +862,10 @@ class MossQwen3ForCausalLM(Qwen3ForCausalLM): batch_size, dtype, device ) for layer_idx in self.deepstack_inject_layer_indices: + # Non-first PP ranks only receive DeepStack payloads for layers + # at or after their local start layer. + if layer_idx < self.model.start_layer: + continue intermediate_tensors[f"deepstack_input_embeds_{layer_idx}"] = torch.zeros( (batch_size, self.config.hidden_size), dtype=dtype, diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 44ce9a88173..07edf2a62d6 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -612,6 +612,7 @@ _SPECULATIVE_DECODING_MODELS = { "DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"), "DeepSeekV4MTPModel": ("vllm.models.deepseek_v4", "DeepSeekV4MTP"), "MiniMaxM3MTP": ("vllm.models.minimax_m3", "MiniMaxM3MTP"), + "BailingMoeV25MTPModel": ("bailing_moe_mtp", "BailingMoeV25MTPModel"), "Gemma4MTPModel": ("gemma4_mtp", "Gemma4MTP"), "ErnieMTPModel": ("ernie_mtp", "ErnieMTP"), "ExaoneMoeMTP": ("exaone_moe_mtp", "ExaoneMoeMTP"), diff --git a/vllm/model_executor/models/transformers/__init__.py b/vllm/model_executor/models/transformers/__init__.py index cb224e5cbc0..78a12876e66 100644 --- a/vllm/model_executor/models/transformers/__init__.py +++ b/vllm/model_executor/models/transformers/__init__.py @@ -16,6 +16,10 @@ # limitations under the License. """Wrapper around `transformers` models""" +from typing import TYPE_CHECKING + +from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS + from vllm.model_executor.models.transformers.base import Base from vllm.model_executor.models.transformers.causal import CausalMixin from vllm.model_executor.models.transformers.legacy import LegacyMixin @@ -32,6 +36,36 @@ from vllm.model_executor.models.transformers.pooling import ( ) from vllm.multimodal import MULTIMODAL_REGISTRY +if TYPE_CHECKING: + import torch + + from vllm.model_executor.layers.attention import Attention + + +def vllm_attention_forward( + # Transformers args + module: "torch.nn.Module", + query: "torch.Tensor", + key: "torch.Tensor", + value: "torch.Tensor", + attention_mask: "torch.Tensor", + # Transformers kwargs + scaling: float | None = None, + # vLLM kwargs + attention_instances: dict[int, "Attention"] | None = None, + **kwargs, +): + self_attn = attention_instances[module.layer_idx] + if scaling is not None: + self_attn.impl.scale = float(scaling) + hidden = query.shape[-2] + query, key, value = (x.transpose(1, 2) for x in (query, key, value)) + query, key, value = (x.reshape(hidden, -1) for x in (query, key, value)) + return self_attn.forward(query, key, value), None + + +ALL_ATTENTION_FUNCTIONS["vllm"] = vllm_attention_forward + # Text only models class TransformersForCausalLM(CausalMixin, Base): ... diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py index bcda62918f3..a36ffacad96 100644 --- a/vllm/model_executor/models/transformers/base.py +++ b/vllm/model_executor/models/transformers/base.py @@ -31,7 +31,6 @@ from transformers.conversion_mapping import ( WeightRenaming, get_model_conversion_mapping, ) -from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS from vllm.compilation.decorators import support_torch_compile from vllm.config.utils import getattr_iter @@ -42,6 +41,7 @@ from vllm.model_executor.layers.attention import ( Attention, EncoderOnlyAttention, ) +from vllm.model_executor.layers.fused_moe import MoERunner from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding from vllm.model_executor.models.interfaces import ( SupportsEagle, @@ -51,6 +51,7 @@ from vllm.model_executor.models.interfaces import ( SupportsQuant, ) from vllm.model_executor.models.interfaces_base import VllmModel +from vllm.model_executor.models.transformers.fuser import BaseFuser, Fusers from vllm.model_executor.models.transformers.utils import ( can_enable_torch_compile, get_feature_request_tip, @@ -58,7 +59,6 @@ from vllm.model_executor.models.transformers.utils import ( log_replacement, replace_conv_class, replace_linear_class, - replace_rms_norm_class, ) from vllm.model_executor.models.utils import ( AutoWeightsLoader, @@ -74,37 +74,10 @@ if TYPE_CHECKING: from transformers import PreTrainedModel from vllm.config import VllmConfig -else: - PreTrainedModel = object logger = init_logger(__name__) -def vllm_flash_attention_forward( - # Transformers args - module: torch.nn.Module, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attention_mask: torch.Tensor, - # Transformers kwargs - scaling: float | None = None, - # vLLM kwargs - attention_instances: dict[int, Attention] | None = None, - **kwargs, -): - self_attn = attention_instances[module.layer_idx] - if scaling is not None: - self_attn.impl.scale = float(scaling) - hidden = query.shape[-2] - query, key, value = (x.transpose(1, 2) for x in (query, key, value)) - query, key, value = (x.reshape(hidden, -1) for x in (query, key, value)) - return self_attn.forward(query, key, value), None - - -ALL_ATTENTION_FUNCTIONS["vllm"] = vllm_flash_attention_forward - - class Base( nn.Module, VllmModel, @@ -141,6 +114,9 @@ class Base( """Ignore unexpected weights whose qualname starts with these prefixes.""" self.ignore_unexpected_suffixes: list[str] = [] """Ignore unexpected weights whose qualname ends with these suffixes.""" + self.packed_modules_mapping: dict[str, list[str]] = {} + """Fused module -> constituent projections, populated by `recursive_replace` + for the quantization machinery and loaders (e.g. bitsandbytes).""" # Attrs for Eagle3 (see self.set_aux_hidden_state_layers) self._target_class: type[nn.Module] = nn.Module @@ -217,7 +193,7 @@ class Base( self.text_config._attn_implementation = "vllm" self.config.dtype = torch.get_default_dtype() - def _get_decoder_cls(self, **kwargs: dict) -> type[PreTrainedModel]: + def _get_decoder_cls(self, **kwargs: dict) -> type["PreTrainedModel"]: """ Get the decoder class from the model. @@ -236,7 +212,7 @@ class Base( def _decorate_cls_for_torch_compile( self, - cls: type[PreTrainedModel], + cls: type["PreTrainedModel"], dynamic_arg_dims: dict[str, int] | None, enable_if: Callable[["VllmConfig"], bool], is_encoder: bool, @@ -356,12 +332,24 @@ class Base( if self.pp_group.world_size <= 1: return - if not self.model.supports_pp_plan: + if self.model.supports_pp_plan: + module = self.model + names = list(module._pp_plan.keys()) + else: + module = self.model.get_decoder() + has_parameters = lambda m: next(m.parameters(), None) is not None + names = [n for n, c in module.named_children() if has_parameters(c)] tip = get_feature_request_tip( self.model_config.model, self.model_config.trust_remote_code ) - raise ValueError( - f"{type(self.model)} does not support pipeline parallel. {tip}" + logger.warning( + "%s does not define a pipeline parallel plan. The Transformers " + "modeling backend will infer the split from the layers of %s in order " + "of declaration and keep parameter-free modules on every rank. This " + "may fail if the model's structure is non-standard. %s", + type(self.model), + type(module), + tip, ) def attrsetter(attr: str) -> Callable[[object, object], None]: @@ -376,10 +364,9 @@ class Base( module_lists = [] module_list_idx = None - pp_plan = list(self.model._pp_plan.keys()) - for i, name in enumerate(pp_plan): + for i, name in enumerate(names): # attrgetter in case the module is nested (e.g. "text_model.layers") - if isinstance(attrgetter(name)(self.model), nn.ModuleList): + if isinstance(attrgetter(name)(module), nn.ModuleList): module_lists.append(name) module_list_idx = i @@ -389,16 +376,16 @@ class Base( "in the base model are not supported yet!" ) if module_list_idx is None: - raise ValueError(f"Could not find `ModuleList` in {type(self.model)}") + raise ValueError(f"Could not find `ModuleList` in {type(module)}") # Layers before module list - for name in pp_plan[:module_list_idx]: + for name in names[:module_list_idx]: if self.pp_group.is_first_rank or ( self._get_tie_word_embeddings() and self.pp_group.is_last_rank ): continue # attrsetter in case the module is nested (e.g. "text_model.embed_tokens") - attrsetter(name)(self.model, PPMissingLayer()) + attrsetter(name)(module, PPMissingLayer()) # Module list start_layer, end_layer = get_pp_indices( @@ -406,41 +393,61 @@ class Base( self.pp_group.rank_in_group, self.pp_group.world_size, ) - layers_name = pp_plan[module_list_idx] + layers_name = names[module_list_idx] # attrgetter in case the module is nested (e.g. "text_model.layers") - layers = attrgetter(layers_name)(self.model) + layers = attrgetter(layers_name)(module) for i in range(len(layers)): if start_layer <= i and i < end_layer: continue layers[i] = PPMissingLayer() # Layers after module list - for name in pp_plan[module_list_idx + 1 :]: + for name in names[module_list_idx + 1 :]: # Modules that should be on last rank if not self.pp_group.is_last_rank: # attrsetter in case the module is nested (e.g. "text_model.norm") - attrsetter(name)(self.model, PPMissingLayer()) + attrsetter(name)(module, PPMissingLayer()) def recursive_replace(self): """Recursively replace modules in the model as needed. Currently, this replaces: + - GLUs with a fused `MergedColumnParallelLinear` + `...AndMul` + - Attention QKV projections with a fused `QKVParallelLinear` + split - `nn.Linear` with vLLM's tensor parallel linear classes - - `*RMSNorm` with vLLM's `RMSNorm` + - `nn.Conv2d` / `nn.Conv3d` with vLLM's `Conv2d` / `Conv3d` + - RMSNorm (detected from their dataflow) with vLLM's `RMSNorm`or `GemmaRMSNorm` """ - tp_plan = self.model.tp_plan + tp_plan = self.model.tp_plan or {} if not tp_plan and self.tp_group.world_size > 1: tip = get_feature_request_tip( self.model_config.model, self.model_config.trust_remote_code ) - raise ValueError( - f"{type(self.model)} does not support tensor parallel. {tip}" + logger.warning_once( + "%s does not define a tensor parallel plan. The Transformers modeling " + "backend will shard the model the best it can during graph fusion and " + "replicate the rest. This may be suboptimal or fail if the model does " + "not fuse cleanly. %s", + type(self.model), + tip, ) # Prefix the patterns because we always start from `self.model` tp_plan = {maybe_prefix("model", k): v for k, v in tp_plan.items()} + # Detect fusable patterns once per module class (cached, so this is cheap) + fusers = Fusers(self.model, self.model_config) + + def register_fusion(fuser: BaseFuser, prefix: str): + """Register a fused layer's mappings just before it is built.""" + orig_to_new_stacked = fuser.orig_to_new_stacked(prefix) + self.hf_to_vllm_mapper.orig_to_new_stacked.update(orig_to_new_stacked) + + packed_modules_mapping = fuser.packed_modules_mapping + self.packed_modules_mapping.update(packed_modules_mapping) + if self.quant_config is not None: + self.quant_config.packed_modules_mapping.update(packed_modules_mapping) def _recursive_replace(module: nn.Module, prefix: str): for child_name, child_module in module.named_children(): @@ -479,11 +486,16 @@ class Base( ) elif isinstance(child_module, (nn.Conv2d, nn.Conv3d)): new_module = replace_conv_class(child_module) - elif child_module.__class__.__name__.endswith("RMSNorm"): - new_module = replace_rms_norm_class( - child_module, self.text_config.hidden_size + elif (fuser := fusers[child_module]) is not None: + register_fusion(fuser, qual_name) + new_module = fuser.fuse( + child_module, qual_name, self.model_config, self.quant_config ) - else: + logger.info_once(fuser.info(child_name)) + _recursive_replace(new_module, prefix=qual_name) + elif not isinstance(child_module, MoERunner): + # MoERunner can contain aliases of shared experts and gates, + # so we don't want to recurse into it and break weight loading. _recursive_replace(child_module, prefix=qual_name) if new_module is not child_module: @@ -538,7 +550,7 @@ class Base( num_heads=num_heads, head_size=head_size, # NOTE: We use Llama scale as default, if it's set by - # Transformers, it's updated in vllm_flash_attention_forward + # Transformers, it's updated in vllm_attention_forward scale=head_size**-0.5, num_kv_heads=num_kv_heads, cache_config=self.cache_config, diff --git a/vllm/model_executor/models/transformers/fuser.py b/vllm/model_executor/models/transformers/fuser.py new file mode 100644 index 00000000000..0aa7c419ceb --- /dev/null +++ b/vllm/model_executor/models/transformers/fuser.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fuser detection for the Transformers modeling backend. + +`get_fuser` traces a module class once (see `fx_utils`) and matches it against +each concrete fuser in `fusers`; `Fusers` caches the result per class for a +whole model. `base.recursive_replace` then applies the matched fuser per +instance. RMSNorm-shaped modules the tracer cannot match are warned about. +""" + +from collections import UserDict +from typing import TYPE_CHECKING + +from cachetools import cached +from torch import nn + +from vllm.logger import init_logger +from vllm.model_executor.models.transformers.fusers import ( + BaseFuser, + GLUFuser, + QKVFuser, + RMSNormFuser, + StackedFuser, +) +from vllm.model_executor.models.transformers.fx_utils import trace + +if TYPE_CHECKING: + from vllm.config.model import ModelConfig + +logger = init_logger(__name__) + + +@cached(cache={}, key=type) +def get_fuser(module: nn.Module) -> BaseFuser | None: + """The fuser for `type(module)` (cached per class), or `None` if no match.""" + # Projection fusions need >=2 sibling linears; the RMSNorm fusion needs a + # leaf module (raw tensor math, no submodules). Nothing else can match, and + # tracing is skipped for it. + n_linear = sum(isinstance(c, nn.Linear) for c in module.children()) + is_leaf = next(module.children(), None) is None + if n_linear < 2 and not is_leaf: + return None + if (graph := trace(module)) is None: + return None + for fuser_cls in (GLUFuser, QKVFuser, RMSNormFuser): + if (fuser := fuser_cls.match(graph, module)) is not None: + if isinstance(fuser, StackedFuser): + try: + fuser.update_forward(module) + except Exception as exc: + # An unrecognised source just means we cannot fuse here. + logger.debug( + "Could not rewrite %s for fusion: %s", type(module), exc + ) + return None + return fuser + # A norm we could not match structurally is left unfused; flag likely misses. + if module.__class__.__name__.endswith("RMSNorm"): + logger.warning_once( + "%s looks like an RMSNorm but its computation did not match the " + "expected pattern, so it was left unfused.", + module.__class__.__name__, + ) + return None + + +class Fusers(UserDict): + """Mapping from module class to fuser, for all fusable classes in a model.""" + + def __init__(self, model: nn.Module, model_config: "ModelConfig"): + self.model_config = model_config + super().__init__({type(m): get_fuser(m) for m in model.modules()}) + + def __getitem__(self, m: nn.Module) -> BaseFuser | None: + fuser = self.data.get(type(m)) + if fuser is not None and fuser.validate(m, self.model_config): + return fuser + return None diff --git a/vllm/model_executor/models/transformers/fusers/__init__.py b/vllm/model_executor/models/transformers/fusers/__init__.py new file mode 100644 index 00000000000..58910b0ecc3 --- /dev/null +++ b/vllm/model_executor/models/transformers/fusers/__init__.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Concrete fusers for the Transformers modeling backend.""" + +from vllm.model_executor.models.transformers.fusers.base import BaseFuser, StackedFuser +from vllm.model_executor.models.transformers.fusers.glu import GLUFuser +from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser +from vllm.model_executor.models.transformers.fusers.qkv import QKVFuser +from vllm.model_executor.models.transformers.fusers.rms_norm import RMSNormFuser + +__all__ = [ + "BaseFuser", + "StackedFuser", + "GLUFuser", + "MoEBlockFuser", + "QKVFuser", + "RMSNormFuser", +] diff --git a/vllm/model_executor/models/transformers/fusers/base.py b/vllm/model_executor/models/transformers/fusers/base.py new file mode 100644 index 00000000000..54fb2d09165 --- /dev/null +++ b/vllm/model_executor/models/transformers/fusers/base.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Base classes for the Transformers backend fusers.""" + +import types +from abc import ABC, abstractmethod +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, ClassVar + +from torch import fx, nn + +from vllm.model_executor.models.utils import ShardId, maybe_prefix + +if TYPE_CHECKING: + from vllm.config.model import ModelConfig + from vllm.model_executor.layers.quantization import QuantizationConfig + + +@dataclass +class BaseFuser(ABC): + """A detected fusion and how to apply it. + + `match` analyses the module *class* once (cached, see `get_fuser`); `fuse` + then applies the fusion to an instance in `recursive_replace`, returning the + module to install in its place. + """ + + @abstractmethod + def info(self, name: str) -> str: + """A human-readable description of the fusion at `name`, for logging.""" + + @classmethod + @abstractmethod + def match(cls, graph: fx.Graph, module: nn.Module) -> "BaseFuser | None": + """Match the pattern in `graph`, returning a fuser if found.""" + + @abstractmethod + def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool: + """Whether this fuser can be applied to this `module` instance.""" + + @abstractmethod + def fuse( + self, + module: nn.Module, + prefix: str, + model_config: "ModelConfig", + quant_config: "QuantizationConfig", + ) -> nn.Module: + """Apply the fusion to an already-validated `module`, returning the + module to install in its place (mutated in place, or freshly built).""" + + def orig_to_new_stacked(self, prefix: str) -> dict[str, tuple[str, ShardId]]: + """`WeightsMapper.orig_to_new_stacked` entries this fuser contributes + (none unless it stacks weights).""" + return {} + + @property + def packed_modules_mapping(self) -> dict[str, list[str]]: + """`packed_modules_mapping` entries this fuser contributes (none unless + it stacks weights).""" + return {} + + +@dataclass +class StackedFuser(BaseFuser): + """A fuser that merges sibling projections into one stacked linear and + rewrites the forward to call it. + + `match` and `update_forward` analyse the class once; `fuse` builds the merged + submodule and binds the compiled forward on an instance in place, so it keeps + its class and any attribute the fusion does not consume. + """ + + merged_name: ClassVar[str] + """Attribute name of the merged module created by `update_attrs`.""" + merged_cls: ClassVar[str] + """Name of the vLLM class the merged projection becomes (for logging).""" + + source_cls: str + """Class of the HF module the fused projections belonged to (for logging).""" + + fused_forward: Callable = field(init=False, repr=False) + """The compiled rewritten forward, set by `update_forward`.""" + + def info(self, name: str) -> str: + sources = " + ".join(shard for shard, _ in self.shards) + return ( + f"Fused: {sources} ({name}: {self.source_cls}) -> " + f"{self.merged_name} ({self.merged_cls})" + ) + + @property + @abstractmethod + def shards(self) -> list[tuple[str, ShardId]]: + """Each projection's original name and its shard id in the merged module. + + Source for both `orig_to_new_stacked` and `packed_modules_mapping`.""" + + def orig_to_new_stacked(self, prefix: str) -> dict[str, tuple[str, ShardId]]: + """`WeightsMapper.orig_to_new_stacked` entries for one fused instance. + + Maps each checkpoint name to `(merged_name, shard_id)`, keyed by qualname + so only this exact layer is remapped, never a same-named projection + elsewhere (e.g. an unfused MoE expert's `gate_proj`).""" + merged = maybe_prefix(prefix, self.merged_name) + return { + maybe_prefix(prefix, name): (merged, shard) for name, shard in self.shards + } + + @property + def packed_modules_mapping(self) -> dict[str, list[str]]: + """`{merged_name: [projection names]}` so quantization can unpack the + fused layer into its per-shard configs.""" + return {self.merged_name: [name for name, _ in self.shards]} + + @abstractmethod + def update_forward(self, module: nn.Module) -> None: + """Rewrite and compile `type(module)`'s forward source. + + Raises if the source does not admit the rewrite (fusion is then skipped). + """ + + @abstractmethod + def update_attrs( + self, + module: nn.Module, + prefix: str, + model_config: "ModelConfig", + quant_config: "QuantizationConfig", + ) -> None: + """Replace `module`'s submodules with the merged module.""" + + def fuse( + self, + module: nn.Module, + prefix: str, + model_config: "ModelConfig", + quant_config: "QuantizationConfig", + ) -> nn.Module: + """Fuse an already-validated `module` in place (see `Fusers.__getitem__`). + + Builds the merged submodule and binds the compiled forward.""" + self.update_attrs(module, prefix, model_config, quant_config) + module.forward = types.MethodType(self.fused_forward, module) + return module diff --git a/vllm/model_executor/models/transformers/fusers/glu.py b/vllm/model_executor/models/transformers/fusers/glu.py new file mode 100644 index 00000000000..951eb35777a --- /dev/null +++ b/vllm/model_executor/models/transformers/fusers/glu.py @@ -0,0 +1,218 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GLU projection fuser: `act(gate(x)) * up(x)` -> a fused gate/up linear.""" + +import ast +import operator +from dataclasses import dataclass +from typing import TYPE_CHECKING, ClassVar + +from torch import fx, nn +from transformers.activations import ACT2CLS + +from vllm.logger import init_logger +from vllm.model_executor.layers.activation import ( + _ACTIVATION_AND_MUL_REGISTRY, + get_act_and_mul_fn, +) +from vllm.model_executor.layers.linear import MergedColumnParallelLinear +from vllm.model_executor.models.transformers.fusers.base import StackedFuser +from vllm.model_executor.models.transformers.fx_utils import ( + compile_forward, + find_node, + is_linear, + peel, + recover_forward, + replace_expr, + single_self_call, +) +from vllm.model_executor.models.transformers.utils import ( + log_replacement, + replace_linear_class, +) +from vllm.model_executor.models.utils import ShardId, maybe_prefix + +if TYPE_CHECKING: + from vllm.config.model import ModelConfig + from vllm.model_executor.layers.quantization import QuantizationConfig + +logger = init_logger(__name__) + + +CLS2ACT: dict[type, list[str]] = {} +for _act_name, _act_cls in ACT2CLS.items(): + if isinstance(_act_cls, tuple): + _act_cls = _act_cls[0] + CLS2ACT.setdefault(_act_cls, []).append(_act_name) + +ACT_AND_MUL_NAMES = frozenset(_ACTIVATION_AND_MUL_REGISTRY.keys()) + + +@dataclass +class GLUFuser(StackedFuser): + """Fuser for the GLU pattern `act(gate(x)) * up(x)`.""" + + act_name: str + gate_name: str + up_name: str + down_name: str | None + merged_name: ClassVar[str] = "gate_up_proj" + merged_cls: ClassVar[str] = "MergedColumnParallelLinear" + + @property + def shards(self) -> list[tuple[str, ShardId]]: + return [(self.gate_name, 0), (self.up_name, 1)] + + @classmethod + def _is_act_of_gate(cls, node: fx.Node, module: nn.Module) -> bool: + """Is node `act(gate(x))` where `gate` is linear and `act` is not linear.""" + return ( + node.op == "call_module" + and not is_linear(node, module) + and len(node.args) == 1 + and isinstance(node.args[0], fx.Node) + and is_linear(node.args[0], module) + ) + + @classmethod + def _get_glu_nodes( + cls, graph: fx.Graph, module: nn.Module + ) -> tuple[fx.Node, fx.Node, fx.Node, fx.Node] | None: + """Search graph for the GLU pattern `act(gate(x)) * up(x)`.""" + for mul in graph.nodes: + if ( + mul.op == "call_function" + and mul.target == operator.mul + and len(mul.args) == 2 + and all(isinstance(arg, fx.Node) for arg in mul.args) + ): + a, b = mul.args + if cls._is_act_of_gate(a, module) and is_linear(b, module): + act, gate, up = a, a.args[0], b + elif cls._is_act_of_gate(b, module) and is_linear(a, module): + act, gate, up = b, b.args[0], a + else: + continue + if ( + all(len(args) == 1 for args in (gate.args, up.args)) + and isinstance(x := gate.args[0], fx.Node) + and x is up.args[0] + ): + return act, gate, up, mul + return None + + @staticmethod + def _get_act_and_mul_name(act: nn.Module) -> str | None: + """Get the name of `act` if it has an `...AndMul` equivalent.""" + for name in CLS2ACT.get(type(act), []): + if name in ACT_AND_MUL_NAMES: + return name + # nn.GELU is not in ACT2CLS, but could be in model code + if type(act) is nn.GELU: + return "gelu_pytorch_tanh" if act.approximate == "tanh" else "gelu" + return None + + @classmethod + def _get_act_and_mul(cls, act: nn.Module) -> nn.Module: + """Get the `...AndMul` equivalent of a Transformers activation module.""" + if name := cls._get_act_and_mul_name(act): + return get_act_and_mul_fn(name) + raise ValueError(f"No AndMul equivalent for {type(act)}") + + @classmethod + def match(cls, graph: fx.Graph, module: nn.Module) -> "GLUFuser | None": + if (glu_nodes := cls._get_glu_nodes(graph, module)) is None: + return None + act_node, gate_node, up_node, mul_node = glu_nodes + + gate = module.get_submodule(gate_node.target) + up = module.get_submodule(up_node.target) + # Shapes must be compatible for a single merged GEMM. + if gate.in_features == up.in_features and (gate.bias is None) == ( + up.bias is None + ): + predicate = lambda n: is_linear(n, module) and peel(n.args[0]) is mul_node + down_node = find_node(graph, predicate) + return cls( + source_cls=type(module).__name__, + act_name=act_node.target, + gate_name=gate_node.target, + up_name=up_node.target, + down_name=down_node.target if down_node is not None else None, + ) + return None + + def update_forward(self, module: nn.Module) -> None: + """Replace `act(gate(x)) * up(x)` with `act(gate_up(x))` in source.""" + funcdef, fn = recover_forward(type(module)) + act_call = single_self_call(funcdef, self.act_name) + gate_call = single_self_call(funcdef, self.gate_name) + up_call = single_self_call(funcdef, self.up_name) + if act_call.args[0] is not gate_call: + raise ValueError("activation does not directly wrap the gate") + if ast.dump(gate_call.args[0]) != ast.dump(up_call.args[0]): + raise ValueError("gate and up inputs are written differently") + muls = [ + node + for node in ast.walk(funcdef) + if isinstance(node, ast.BinOp) + and isinstance(node.op, ast.Mult) + and {id(node.left), id(node.right)} == {id(act_call), id(up_call)} + ] + if len(muls) != 1: + raise ValueError("no multiply of the activation and up projection") + + # act(gate(x)) * up(x) -> act(gate_up(x)) + assert isinstance(gate_call.func, ast.Attribute) + gate_call.func.attr = self.merged_name + replace_expr(funcdef, muls[0], act_call) + self.fused_forward = compile_forward(funcdef, fn) + + def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool: + act = module.get_submodule(self.act_name) + if self._get_act_and_mul_name(act) is None: + logger.debug("No AndMul equivalent for %s; skipping fusion", type(act)) + return False + return True + + def update_attrs( + self, + module: nn.Module, + prefix: str, + model_config: "ModelConfig", + quant_config: "QuantizationConfig", + ) -> None: + act_fn = self._get_act_and_mul(module.get_submodule(self.act_name)) + gate = module.get_submodule(self.gate_name) + up = module.get_submodule(self.up_name) + merged = MergedColumnParallelLinear( + input_size=gate.in_features, + output_sizes=[gate.out_features, up.out_features], + bias=gate.bias is not None, + quant_config=quant_config, + prefix=maybe_prefix(prefix, self.merged_name), + return_bias=False, + ) + logger.debug( + "%s: %s, %s: %s -> %s: %s", + self.gate_name, + gate, + self.up_name, + up, + self.merged_name, + merged, + ) + setattr(module, self.merged_name, merged) + setattr(module, self.act_name, act_fn) + # Drop the consumed submodules so their (meta) params are not expected. + delattr(module, self.gate_name) + delattr(module, self.up_name) + # If there is a down projection, we know it must be rowwise. + if self.down_name is not None: + down_prefix = maybe_prefix(prefix, self.down_name) + down = module.get_submodule(self.down_name) + new_down = replace_linear_class( + down, "rowwise", quant_config, prefix=down_prefix + ) + setattr(module, self.down_name, new_down) + log_replacement(down_prefix, down, new_down) diff --git a/vllm/model_executor/models/transformers/fusers/moe.py b/vllm/model_executor/models/transformers/fusers/moe.py new file mode 100644 index 00000000000..6a3c7e85d6e --- /dev/null +++ b/vllm/model_executor/models/transformers/fusers/moe.py @@ -0,0 +1,268 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MoE fuser: route an HF MoE block through `FusedMoE` with vLLM's own routing.""" + +import ast +import inspect +import textwrap +import types +from collections.abc import Iterator +from dataclasses import dataclass +from itertools import chain + +import torch +from torch import fx, nn + +from vllm.distributed import tensor_model_parallel_all_gather +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.models.transformers.fx_utils import ( + find_node, + is_op, + peel, + trace, +) +from vllm.model_executor.models.utils import maybe_prefix, sequence_parallel_chunk + + +def named_state(module: nn.Module) -> Iterator[tuple[str, torch.Tensor]]: + """`module`'s own state (i.e. named parameters and buffers).""" + return chain(module.named_parameters(), module.named_buffers()) + + +def _own_returns(node: ast.AST) -> Iterator[ast.Return]: + """`return` statements in `node`'s own scope, not in nested functions.""" + stack = list(ast.iter_child_nodes(node)) + while stack: + child = stack.pop() + if isinstance(child, ast.Return): + yield child + elif not isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)): + stack.extend(ast.iter_child_nodes(child)) + + +def _returns_tuple(cls: type[nn.Module]) -> bool: + """Does `cls.forward()` return a tuple?""" + try: + source = textwrap.dedent(inspect.getsource(inspect.unwrap(cls.forward))) + forward = ast.parse(source).body[0] + except (OSError, SyntaxError, TypeError, IndexError): + return True + # Names bound to a tuple literal, e.g. `out = hidden, logits` then `return out`. + tuple_names = { + target.id + for node in ast.walk(forward) + if isinstance(node, ast.Assign) and isinstance(node.value, ast.Tuple) + for target in node.targets + if isinstance(target, ast.Name) + } + + def yields_tuple(value: ast.expr | None) -> bool: + if isinstance(value, ast.Tuple): + return True + if isinstance(value, ast.Name): + return value.id in tuple_names + if isinstance(value, ast.IfExp): + return yields_tuple(value.body) or yields_tuple(value.orelse) + return False + + return any(yields_tuple(ret.value) for ret in _own_returns(forward)) + + +def _is_scalar_gate(module: nn.Module) -> bool: + """A linear projecting to a single logit (the shared-expert sigmoid gate).""" + weight = getattr(module, "weight", None) + return ( + isinstance(module, nn.Linear) + and weight is not None + and weight.ndim == 2 + and weight.shape[0] == 1 + ) + + +def _reaches(node: fx.Node, key: str) -> set[fx.Node]: + """Returns the set of nodes reachable from `node` by following `key` edges.""" + seen: set[fx.Node] = set() + stack = [node] + while stack: + n = stack.pop() + if n in seen: + continue + seen.add(n) + stack.extend(getattr(n, key)) + return seen + + +class SharedExpertMLP(nn.Module): + """Wraps an HF shared expert, applying the output gating it is paired with.""" + + def __init__(self, shared_experts: nn.Module, gate: nn.Module | None = None): + super().__init__() + self.shared_experts = shared_experts + self.gate = gate + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + out = self.shared_experts(hidden_states) + if self.gate is not None: + out = torch.sigmoid(self.gate(hidden_states)[0]) * out + return out + + +def _moe_block_forward(self: nn.Module, hidden_states: torch.Tensor) -> torch.Tensor: + """Standard MoE block forward. + + Routing and any shared experts are handled inside `self.experts: MoERunner`.""" + orig_shape = hidden_states.shape + hidden_states = hidden_states.reshape(-1, orig_shape[-1]) + num_tokens = hidden_states.shape[0] + is_sequence_parallel = self.experts.moe_config.is_sequence_parallel + if is_sequence_parallel: + hidden_states = sequence_parallel_chunk(hidden_states) + out = self.experts(hidden_states, router_logits=hidden_states) + if is_sequence_parallel: + out = tensor_model_parallel_all_gather(out, 0)[:num_tokens] + return out.reshape(orig_shape) + + +@dataclass +class MoEBlockFuser: + """Fuser for MoE block `experts`, `gate` and `shared_experts` (optional).""" + + gate_name: str + scoring_func: str + shared_name: str | None + shared_gate_name: str | None + + @staticmethod + def _match_router(gate: nn.Module) -> str | None: + """Matches `topk(score(linear(x)))`, `score` being `softmax`/`sigmoid`.""" + if [name for name, _ in named_state(gate)] != ["weight"]: + return None + graph = trace(gate) + if graph is None: + return None + topk = find_node(graph, lambda n: is_op(n, "topk")) + if topk is None: + return None + # Exactly one scoring op upstream of the top-k, fed (transitively) by a linear. + scorers = [ + n + for n in _reaches(topk, "all_input_nodes") + if is_op(n, "softmax") or is_op(n, "sigmoid") + ] + if len(scorers) != 1: + return None + scorer = scorers[0] + if not any(is_op(n, "linear") for n in _reaches(scorer, "all_input_nodes")): + return None + return "softmax" if is_op(scorer, "softmax") else "sigmoid" + + @staticmethod + def _match_shared_experts( + graph: fx.Graph, experts: str + ) -> tuple[str | None, str | None]: + """Detects the shared expert and its optional gate by dataflow.""" + experts_predicate = lambda n: n.op == "call_module" and n.target == experts + if (experts_node := find_node(graph, experts_predicate)) is None: + return None, None + from_experts = _reaches(experts_node, "users") + for add in graph.nodes: + if not is_op(add, "add"): + continue + operands = [a for a in add.args if isinstance(a, fx.Node)] + # Exactly one side is the experts' output; the other is the shared path. + sides = [a in from_experts for a in operands] + if len(operands) != 2 or sides.count(True) != 1: + continue + cone = _reaches(operands[sides.index(False)], "all_input_nodes") + modules = [n for n in cone if n.op == "call_module" and n.target != experts] + # A sigmoid wrapping one of those modules marks the shared-expert gate. + gate = next( + ( + src + for n in cone + if is_op(n, "sigmoid") + and isinstance(src := peel(n.args[0]), fx.Node) + and src in modules + ), + None, + ) + shared = [n for n in modules if n is not gate] + if len(shared) != 1: + return None, None + return shared[0].target, (gate.target if gate is not None else None) + return None, None + + @classmethod + def match(cls, moe_block: nn.Module, experts_name: str) -> "MoEBlockFuser | None": + # Standard MoE block returns a single tensor. + if _returns_tuple(type(moe_block)): + return None + # Router: the child that scores + top-k selects. + gate_name = scoring_func = None + for name, child in moe_block.named_children(): + if name != experts_name and (func := cls._match_router(child)) is not None: + gate_name, scoring_func = name, func + break + if gate_name is None or scoring_func is None: + return None + # Shared expert: a child the block adds to the experts' output. + shared_name = shared_gate_name = None + others = [ + n + for n, _ in moe_block.named_children() + if n not in {experts_name, gate_name} + ] + if others: + graph = trace(moe_block) + if graph is None: + return None + shared_name, shared_gate_name = cls._match_shared_experts( + graph, experts_name + ) + if shared_gate_name is not None and not _is_scalar_gate( + getattr(moe_block, shared_gate_name) + ): + return None + # Fail closed: `rewrite_forward` runs only the experts and the detected + # shared expert, so any other stateful child would be dropped. + accounted = {experts_name, gate_name, shared_name, shared_gate_name} + for name, child in moe_block.named_children(): + if name not in accounted and next(named_state(child), None) is not None: + return None + return cls(gate_name, scoring_func, shared_name, shared_gate_name) + + def gate(self, moe_block: nn.Module, prefix: str) -> ReplicatedLinear: + """Rebuild the HF gate as a `ReplicatedLinear` for vLLM's fused MoE.""" + num_experts, hidden_size = getattr(moe_block, self.gate_name).weight.shape + gate = ReplicatedLinear( + hidden_size, + num_experts, + bias=False, + prefix=maybe_prefix(prefix, self.gate_name), + ) + setattr(moe_block, self.gate_name, gate) + return gate + + def shared_experts( + self, moe_block: nn.Module, prefix: str + ) -> SharedExpertMLP | None: + """Build the HF shared expert (and its optional gate) + as a `SharedExpertMLP` for vLLM's fused MoE.""" + if self.shared_name is None: + return None + shared_experts = getattr(moe_block, self.shared_name) + gate = None + if self.shared_gate_name is not None: + hf_gate = getattr(moe_block, self.shared_gate_name) + gate = ReplicatedLinear( + hf_gate.in_features, + hf_gate.out_features, + bias=hf_gate.bias is not None, + prefix=maybe_prefix(prefix, self.shared_gate_name), + ) + setattr(moe_block, self.shared_gate_name, gate) + return SharedExpertMLP(shared_experts, gate) + + def rewrite_forward(self, moe_block: nn.Module) -> None: + """Rewrite `moe_block.forward` to route through vLLM's fused MoE.""" + moe_block.forward = types.MethodType(_moe_block_forward, moe_block) diff --git a/vllm/model_executor/models/transformers/fusers/qkv.py b/vllm/model_executor/models/transformers/fusers/qkv.py new file mode 100644 index 00000000000..03f5a0b49e4 --- /dev/null +++ b/vllm/model_executor/models/transformers/fusers/qkv.py @@ -0,0 +1,220 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""QKV projection fuser: `q(x), k(x), v(x)` -> a fused qkv linear + split.""" + +import ast +from dataclasses import dataclass +from typing import TYPE_CHECKING, ClassVar + +from torch import fx, nn + +from vllm.logger import init_logger +from vllm.model_executor.layers.linear import QKVParallelLinear +from vllm.model_executor.models.transformers.fusers.base import StackedFuser +from vllm.model_executor.models.transformers.fx_utils import ( + compile_forward, + innermost_block, + is_linear, + recover_forward, + replace_expr, + single_self_call, +) +from vllm.model_executor.models.transformers.utils import ( + log_replacement, + replace_linear_class, +) +from vllm.model_executor.models.utils import ShardId, maybe_prefix + +if TYPE_CHECKING: + from vllm.config.model import ModelConfig + from vllm.model_executor.layers.quantization import QuantizationConfig + +logger = init_logger(__name__) + + +@dataclass +class QKVFuser(StackedFuser): + """Fuser for the attention QKV pattern `q(x), k(x), v(x)`.""" + + q_name: str + k_name: str + v_name: str + o_name: str | None + merged_name: ClassVar[str] = "qkv_proj" + merged_cls: ClassVar[str] = "QKVParallelLinear" + + @property + def shards(self) -> list[tuple[str, ShardId]]: + return [(self.q_name, "q"), (self.k_name, "k"), (self.v_name, "v")] + + @classmethod + def _get_qkv_nodes( + cls, graph: fx.Graph, module: nn.Module + ) -> tuple[fx.Node, fx.Node, fx.Node] | None: + """Search `graph` for the QKV pattern `q(x), k(x), v(x)`.""" + by_input: dict[fx.Node, list[fx.Node]] = {} + for node in graph.nodes: + if ( + is_linear(node, module) + and len(node.args) == 1 + and not node.kwargs + and isinstance(node.args[0], fx.Node) + and node.args[0].op == "placeholder" + ): + by_input.setdefault(node.args[0], []).append(node) + triples = [nodes for nodes in by_input.values() if len(nodes) == 3] + if len(triples) != 1: + return None + + q_node, k_node, v_node = nodes = triples[0] + outs = [module.get_submodule(node.target).out_features for node in nodes] + if len(set(outs)) == 2: + # q is identified as the larger projection (GQA) + (q_node,) = (n for n, out in zip(nodes, outs) if outs.count(out) == 1) + k_node, v_node = (n for n, out in zip(nodes, outs) if outs.count(out) == 2) + if module.get_submodule(q_node.target).out_features != max(outs): + return None + elif len(set(outs)) != 1: + return None + return q_node, k_node, v_node + + @classmethod + def match(cls, graph: fx.Graph, module: nn.Module) -> "QKVFuser | None": + if (qkv_nodes := cls._get_qkv_nodes(graph, module)) is None: + return None + q, k, v = qkv_nodes + names = dict(q_name=q.target, k_name=k.target, v_name=v.target) + attn_width = module.get_submodule(q.target).out_features + candidates = [ + name + for name, child in module.named_children() + if isinstance(child, nn.Linear) + and name not in names.values() + and child.in_features == attn_width + ] + names["o_name"] = candidates[0] if len(candidates) == 1 else None + return cls(source_cls=type(module).__name__, **names) + + def update_forward(self, module: nn.Module) -> None: + """Replace `q(x), k(x), v(x)` with `qkv(x).split(sizes, -1)` in source.""" + funcdef, fn = recover_forward(type(module)) + calls = [ + single_self_call(funcdef, name) + for name in (self.q_name, self.k_name, self.v_name) + ] + arg_dumps = {ast.dump(call.args[0]) for call in calls} + if len(arg_dumps) != 1: + raise ValueError("projection inputs are written differently") + # The trace may be partial, so prove projection exclusivity in source: + # no other linear child may consume the same input (else the matched + # three may not be q, k and v) + other_linears = { + name + for name, child in module.named_children() + if isinstance(child, nn.Linear) + } - {self.q_name, self.k_name, self.v_name} + for node in ast.walk(funcdef): + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr in other_linears + and any(ast.dump(arg) in arg_dumps for arg in node.args) + ): + raise ValueError("another linear consumes the same input") + blocks = [innermost_block(funcdef.body, call) for call in calls] + if any(found is None for found in blocks): + raise ValueError("projection calls not found in the function body") + if len({id(block) for block, _ in blocks}) != 1: + raise ValueError("projection calls are in different blocks") + + # q(x), k(x), v(x) -> q, k, v = qkv(x).split(self.qkv.split_sizes, -1) + names = {node.id for node in ast.walk(funcdef) if isinstance(node, ast.Name)} + temps = [f"{name}_fused" for name in (self.q_name, self.k_name, self.v_name)] + if names & set(temps): + raise ValueError("fused temporaries would shadow existing names") + merged = f"self.{self.merged_name}" + template = ( + f"{', '.join(temps)} = {merged}(__arg__).split({merged}.split_sizes, -1)" + ) + assign = ast.parse(template).body[0] + arg = next( + node + for node in ast.walk(assign) + if isinstance(node, ast.Name) and node.id == "__arg__" + ) + replace_expr(assign, arg, calls[0].args[0]) + block, index = blocks[0] + ast.copy_location(assign, block[index]) + block.insert(min(index for _, index in blocks), assign) + for call, temp in zip(calls, temps): + replace_expr(funcdef, call, ast.Name(id=temp, ctx=ast.Load())) + self.fused_forward = compile_forward(funcdef, fn) + + def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool: + """Shapes must be compatible for a single merged, head-sharded GEMM.""" + q = module.get_submodule(self.q_name) + k = module.get_submodule(self.k_name) + v = module.get_submodule(self.v_name) + head_size = model_config.get_head_size() + compatible = ( + q.in_features == k.in_features == v.in_features + and len({proj.bias is None for proj in (q, k, v)}) == 1 + and k.out_features == v.out_features + and q.out_features % head_size == 0 + and k.out_features % head_size == 0 + ) + if not compatible: + logger.debug("%s is not compatible with QKV fusion", type(module)) + return compatible + + def update_attrs( + self, + module: nn.Module, + prefix: str, + model_config: "ModelConfig", + quant_config: "QuantizationConfig", + ) -> None: + head_size = model_config.get_head_size() + q = module.get_submodule(self.q_name) + k = module.get_submodule(self.k_name) + merged = QKVParallelLinear( + hidden_size=q.in_features, + head_size=head_size, + total_num_heads=q.out_features // head_size, + total_num_kv_heads=k.out_features // head_size, + bias=q.bias is not None, + quant_config=quant_config, + prefix=maybe_prefix(prefix, self.merged_name), + return_bias=False, + ) + logger.debug( + "%s: %s, %s: %s, %s: %s -> %s: %s", + self.q_name, + q, + self.k_name, + k, + self.v_name, + module.get_submodule(self.v_name), + self.merged_name, + merged, + ) + # The rewritten forward splits the merged projection into this rank's + # shard sizes (see `update_forward`) + merged.split_sizes = [ + merged.num_heads * merged.head_size, + merged.num_kv_heads * merged.head_size, + merged.num_kv_heads * merged.v_head_size, + ] + setattr(module, self.merged_name, merged) + # Drop the consumed submodules so their (meta) params are not expected. + for name in (self.q_name, self.k_name, self.v_name): + delattr(module, name) + # If there is an output projection, we know it must be rowwise. + if self.o_name is not None: + o_proj_prefix = maybe_prefix(prefix, self.o_name) + o_proj = module.get_submodule(self.o_name) + new_o = replace_linear_class( + o_proj, "rowwise", quant_config, prefix=o_proj_prefix + ) + setattr(module, self.o_name, new_o) + log_replacement(o_proj_prefix, o_proj, new_o) diff --git a/vllm/model_executor/models/transformers/fusers/rms_norm.py b/vllm/model_executor/models/transformers/fusers/rms_norm.py new file mode 100644 index 00000000000..829fd9a541c --- /dev/null +++ b/vllm/model_executor/models/transformers/fusers/rms_norm.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""RMSNorm fuser: detect the norm structurally and swap in vLLM's fused RMSNorm.""" + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch +from torch import fx, nn + +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_gather, +) +from vllm.distributed.parallel_state import model_parallel_is_initialized +from vllm.distributed.utils import split_tensor_along_last_dim +from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm +from vllm.model_executor.models.transformers.fusers.base import BaseFuser +from vllm.model_executor.models.transformers.fx_utils import ( + find_node, + forward_input_count, + is_op, + peel, + trace, +) + +if TYPE_CHECKING: + from vllm.config.model import ModelConfig + from vllm.model_executor.layers.quantization import QuantizationConfig + + +def _is_squared(node: object, x: fx.Node) -> bool: + """`x**2`, `x.square()` or `x * x`, through any dtype casts.""" + node = peel(node) + if is_op(node, "pow"): + base, exp = node.args + return peel(base) is x and exp == 2 + if is_op(node, "square"): + return peel(node.args[0]) is x + if is_op(node, "mul"): + a, b = node.args + return peel(a) is x and peel(b) is x + return False + + +def _variance_eps(rsqrt: fx.Node, x: fx.Node) -> float | None: + """eps from `rsqrt(mean(x**2, -1) + eps)`, or `None` if not that shape.""" + add = peel(rsqrt.args[0]) + if not is_op(add, "add"): + return None + consts = [a for a in add.args if isinstance(a, (int, float))] + nodes = [a for a in add.args if isinstance(a, fx.Node)] + if len(consts) != 1 or len(nodes) != 1: + return None + mean = peel(nodes[0]) + if not is_op(mean, "mean"): + return None + if not _is_squared(mean.args[0], x): + return None + return float(consts[0]) + + +def _is_one_plus(node: object) -> bool: + """`1 + weight` in either operand order (marks a zero-centered weight).""" + node = peel(node) + if not is_op(node, "add"): + return False + return any(isinstance(a, (int, float)) and a == 1 for a in node.args) + + +def _has_trailing_compute(graph: fx.Graph, node: fx.Node) -> bool: + """Does the forward compute anything after `node` before returning?""" + output = find_node(graph, lambda n: n.op == "output") + if output is None or not output.args: + return False + return peel(output.args[0]) is not node + + +class TPAwareNormMixin(nn.Module): + """Mixin for RMSNorms that reconstructs a TP-sharded input before normalizing.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + if model_parallel_is_initialized(): + self.tp_size = get_tensor_model_parallel_world_size() + self.tp_rank = get_tensor_model_parallel_rank() + else: + self.tp_size, self.tp_rank = 1, 0 + + def forward( + self, x: torch.Tensor, residual: torch.Tensor | None = None + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if self.tp_size > 1 and x.shape[-1] < (full := self.weight.shape[0]): + if x.shape[-1] * self.tp_size != full: + raise ValueError( + f"Cannot gather norm of width {full}: a TP-sharded input of " + f"width {x.shape[-1]} does not tile it evenly across " + f"{self.tp_size} ranks (replicated or uneven sharding)." + ) + x = tensor_model_parallel_all_gather(x.contiguous()) + x = super().forward(x) + splits = split_tensor_along_last_dim(x, num_partitions=self.tp_size) + return splits[self.tp_rank] + return super().forward(x, residual) + + +class TPAwareRMSNorm(TPAwareNormMixin, RMSNorm): + """`RMSNorm` that reconstructs a TP-sharded input before normalizing.""" + + +class TPAwareGemmaRMSNorm(TPAwareNormMixin, GemmaRMSNorm): + """`GemmaRMSNorm` that reconstructs a TP-sharded input before normalizing.""" + + +@dataclass +class RMSNormFuser(BaseFuser): + """Fuser for RMSNorm patterns, including Gemma-style zero-centered weights.""" + + zero_centered: bool + """Gemma-style `(1 + weight)` scaling (weight initialised at zero).""" + source_cls: str + """Class name of the norm this was matched from (for logging).""" + + def info(self, name: str) -> str: + norm = "GemmaRMSNorm" if self.zero_centered else "RMSNorm" + return f"Fused: {name} ({self.source_cls}) -> {norm} (CustomOp)" + + @classmethod + def match(cls, graph: fx.Graph, module: nn.Module) -> "RMSNormFuser | None": + """Match a graph to the RMSNorm pattern, returning a fuser if found.""" + if forward_input_count(type(module)) != 1: + return None + x = find_node(graph, lambda n: n.op == "placeholder") + if x is None: + return None + # Handle native torch `rms_norm` op. + rms_norm = find_node(graph, lambda n: is_op(n, "rms_norm")) + if rms_norm is not None and rms_norm.args and peel(rms_norm.args[0]) is x: + if _has_trailing_compute(graph, rms_norm): + return None + return cls(zero_centered=False, source_cls=type(module).__name__) + # Handle explicit `x * rsqrt(mean(x**2, -1) + eps)` pattern. + # The rsqrt over the mean-square variance is the spine of the norm. + rsqrt = None + for node in graph.nodes: + if is_op(node, "rsqrt") and _variance_eps(node, x) is not None: + rsqrt = node + break + if rsqrt is None: + return None + # The `x * rsqrt(...)` normalize multiply. + normalize = find_node( + graph, lambda n: is_op(n, "mul") and rsqrt in map(peel, n.args) + ) + if normalize is None: + return None + # An optional later `weight * normalized` (or `(1 + weight) * normalized`). + tail, zero_centered = normalize, False + for node in graph.nodes: + if not is_op(node, "mul") or node is normalize: + continue + operands = [peel(a) for a in node.args if isinstance(a, fx.Node)] + if len(operands) == 2 and normalize in operands: + weight = next(o for o in operands if o is not normalize) + tail, zero_centered = node, _is_one_plus(weight) + break + # The norm must be the last compute in forward, or it is not a pure norm. + if _has_trailing_compute(graph, tail): + return None + return cls(zero_centered=zero_centered, source_cls=type(module).__name__) + + @staticmethod + def _eps_from_graph(graph: fx.Graph) -> float | None: + """Extract the `eps` constant from the graph, if present.""" + if (x := find_node(graph, lambda n: n.op == "placeholder")) is None: + return None + fused = find_node(graph, lambda n: is_op(n, "rms_norm")) + if fused is not None and fused.args and peel(fused.args[0]) is x: + args, kwargs = fused.args, fused.kwargs + eps = args[3] if len(args) > 3 else kwargs.get("eps") + return eps if isinstance(eps, (int, float)) else None + for node in graph.nodes: + if is_op(node, "rsqrt") and (eps := _variance_eps(node, x)) is not None: + return eps + return None + + def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool: + return True + + def fuse( + self, + module: nn.Module, + prefix: str, + model_config: "ModelConfig", + quant_config: "QuantizationConfig", + ) -> nn.Module: + """Fuse the matched RMSNorm pattern into a vLLM fused RMSNorm CustomOp.""" + weight = getattr(module, "weight", None) + hidden_size = ( + weight.size(0) if weight is not None else model_config.get_hidden_size() + ) + graph = trace(module) + eps = self._eps_from_graph(graph) if graph is not None else None + if eps is None: + # If eps not in graph, match torch behaviour. + dtype = weight.dtype if weight is not None else model_config.dtype + eps = torch.finfo(dtype).eps + if self.zero_centered: + return TPAwareGemmaRMSNorm(hidden_size=hidden_size, eps=eps) + has_weight = weight is not None + return TPAwareRMSNorm( + hidden_size=hidden_size, + eps=eps, + has_weight=has_weight, + dtype=weight.dtype if has_weight else None, + ) diff --git a/vllm/model_executor/models/transformers/fx_utils.py b/vllm/model_executor/models/transformers/fx_utils.py new file mode 100644 index 00000000000..0e043941d8a --- /dev/null +++ b/vllm/model_executor/models/transformers/fx_utils.py @@ -0,0 +1,273 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""fx tracing and forward-source rewriting for the Transformers backend fusers. + +A small engine, independent of any particular pattern: trace a module's forward +with `torch.fx` (tolerating a partial graph), inspect the resulting nodes, and +rewrite the forward's *source* (AST) so only matched calls change while the rest +stays live Python. `fusion.py` builds the concrete fusion patterns on top. +""" + +import ast +import inspect +import operator +import textwrap +from collections.abc import Callable + +import torch +from torch import fx, nn +from torch.nn import functional as F + +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +def _infer_len(node: fx.Node) -> int | None: + """Concrete length of a proxy's value, inferred from its node chain. + + Lets tracing pass through the shape unpacks and `*`-splats (e.g. + `(*input_shape, -1, head_dim)`) that precede the patterns in HF attention. + """ + # `x.shape` has the rank of `x`, when known + if ( + node.op == "call_function" + and node.target is getattr + and node.args[1] == "shape" + and (rank := _rank(node.args[0])) is not None + ): + return rank + # Slices of known-length values + if node.op == "call_function" and node.target is operator.getitem: + src_len = _infer_len(node.args[0]) + index = node.args[1] + if src_len is not None and isinstance(index, slice): + return len(range(*index.indices(src_len))) + return None + + +def _rank(node: fx.Node) -> int | None: + """The tensor rank of `node`'s value, if known.""" + # vLLM always feeds the model [1, seq_len, hidden_size] hidden states + if node.op == "placeholder" and node.target == "hidden_states": + return 3 + return None + + +class _SizedProxy(fx.Proxy): + """Proxy whose `len` is inferred from the graph (see `_infer_len`).""" + + def __len__(self) -> int: + length = _infer_len(self.node) + if length is None: + return super().__len__() + return length + + +class _AllLeafTracer(fx.Tracer): + """Tracer that treats every submodule as a leaf. + + Each child stays one `call_module` node, so matching sees the module's own + forward structure (activations aren't decomposed into e.g. `sigmoid * x`). + `iter` traces through the leading shape unpacks (see `_infer_len`); anything + else untraceable ends the trace early and the partial graph is matched. + """ + + def is_leaf_module(self, m: nn.Module, module_qualified_name: str) -> bool: + return True + + def proxy(self, node: fx.Node) -> fx.Proxy: + return _SizedProxy(node, self) + + def iter(self, obj: fx.Proxy): + length = _infer_len(obj.node) + if length is None: + return super().iter(obj) + return iter([obj[i] for i in range(length)]) + + +def trace(module: nn.Module) -> fx.Graph | None: + """Trace `module.forward`, returning the partial graph on failure. + + The graph is only evidence for matching, and the patterns sit at the top of + their forwards, so a trace that fails partway can still be matched.""" + tracer = _AllLeafTracer() + try: + return tracer.trace(module) + except Exception as exc: + logger.debug("Could not fully trace %s: %s", type(module), exc) + return getattr(tracer, "graph", None) + + +def recover_forward(cls: type[nn.Module]) -> tuple[ast.FunctionDef, Callable]: + """Parse the source of `cls.forward`, ready for rewriting.""" + fn = inspect.unwrap(cls.forward) + if fn.__code__.co_freevars: + raise ValueError("forward is a closure") + tree = ast.parse(textwrap.dedent(inspect.getsource(fn))) + funcdef = tree.body[0] + if not isinstance(funcdef, ast.FunctionDef): + raise ValueError("source is not a plain function definition") + # `fn` is already unwrapped; don't re-apply its decorators + funcdef.decorator_list.clear() + # Annotations may not evaluate outside the defining module (e.g. with + # postponed evaluation); they're not needed at runtime + funcdef.returns = None + args = funcdef.args + for arg in ( + *args.posonlyargs, + *args.args, + *args.kwonlyargs, + *filter(None, (args.vararg, args.kwarg)), + ): + arg.annotation = None + # Recompiling outside the class body would break name mangling + for node in ast.walk(funcdef): + name = getattr(node, "attr", None) or getattr(node, "id", None) + if name and name.startswith("__") and not name.endswith("__"): + raise ValueError(f"{name} would be name mangled") + return funcdef, fn + + +def forward_input_count(cls: type[nn.Module]) -> int: + """The number of tensor inputs `cls.forward` declares, excluding `self` and + any `*args`/`**kwargs`. Read from the signature, so it is independent of + whether the trace completes (unlike counting placeholders).""" + try: + params = list(inspect.signature(cls.forward).parameters.values())[1:] + except (ValueError, TypeError): + return 1 # uninspectable: assume a single input and let matching decide + fixed = ( + inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + ) + return sum(1 for p in params if p.kind in fixed) + + +def compile_forward(funcdef: ast.FunctionDef, fn: Callable) -> Callable: + """Compile `funcdef` in `fn`'s module so tracebacks point at the source.""" + module = ast.Module(body=[funcdef], type_ignores=[]) + ast.fix_missing_locations(module) + ast.increment_lineno(module, fn.__code__.co_firstlineno - 1) + code = compile(module, fn.__code__.co_filename, "exec") + namespace: dict = {} + exec(code, fn.__globals__, namespace) + return namespace[funcdef.name] + + +def single_self_call(funcdef: ast.FunctionDef, name: str) -> ast.Call: + """The unique `self.(arg)` call in `funcdef`. + + Raises unless `name` appears exactly once, as such a call, so the source + rewrite agrees with the fx match. + """ + uses = [ + node + for node in ast.walk(funcdef) + if isinstance(node, ast.Attribute) and node.attr == name + ] + if len(uses) != 1: + raise ValueError(f"{name} is referenced {len(uses)} times") + calls = [ + node + for node in ast.walk(funcdef) + if isinstance(node, ast.Call) + and node.func is uses[0] + and len(node.args) == 1 + and not isinstance(node.args[0], ast.Starred) + and not node.keywords + ] + if ( + len(calls) != 1 + or not isinstance(uses[0].value, ast.Name) + or uses[0].value.id != "self" + ): + raise ValueError(f"{name} is not a single-argument call on self") + return calls[0] + + +def innermost_block( + block: list[ast.stmt], node: ast.AST +) -> tuple[list[ast.stmt], int] | None: + """The innermost statement list containing `node`, and the index within.""" + for index, stmt in enumerate(block): + if not any(child is node for child in ast.walk(stmt)): + continue + child_blocks = [ + getattr(stmt, fld, None) for fld in ("body", "orelse", "finalbody") + ] + child_blocks += [h.body for h in getattr(stmt, "handlers", [])] + child_blocks += [c.body for c in getattr(stmt, "cases", [])] + for child_block in child_blocks: + if ( + isinstance(child_block, list) + and child_block + and (found := innermost_block(child_block, node)) is not None + ): + return found + return block, index + return None + + +def replace_expr(module: ast.AST, old: ast.expr, new: ast.expr) -> None: + """Replace the expression `old` (by identity) with `new` within `module`.""" + + class _Replacer(ast.NodeTransformer): + def visit(self, node: ast.AST) -> ast.AST: + if node is old: + return new + return super().generic_visit(node) + + _Replacer().visit(module) + + +def find_node(graph: fx.Graph, predicate: Callable[[fx.Node], bool]) -> fx.Node | None: + """The first node in `graph` matching `predicate`, or `None`.""" + return next((n for n in graph.nodes if predicate(n)), None) + + +def is_linear(node: fx.Node, module: nn.Module) -> bool: + """Is node `nn.Linear.__call__()`.""" + return node.op == "call_module" and isinstance( + module.get_submodule(node.target), nn.Linear + ) + + +_DTYPE_CASTS = frozenset({"to", "float", "double", "half", "bfloat16", "type_as"}) + + +def peel(node: object) -> object: + """Strip dtype-cast wrappers (`.to(...)`, `.float()`, `.type_as(...)`).""" + while ( + isinstance(node, fx.Node) + and node.op == "call_method" + and node.target in _DTYPE_CASTS + ): + node = node.args[0] + return node + + +def is_fn(node: object, target: Callable) -> bool: + """Is node `()`.""" + return ( + isinstance(node, fx.Node) + and node.op == "call_function" + and node.target is target + ) + + +def is_method(node: object, name: str) -> bool: + """Is node `.()`.""" + return ( + isinstance(node, fx.Node) and node.op == "call_method" and node.target == name + ) + + +def is_op(node: object, name: str) -> bool: + """ + Is node `torch.()`, `F.()`, `operator.()`, or `Tensor.()`. + """ + return any( + is_fn(node, getattr(module, name, None)) for module in (torch, F, operator) + ) or (hasattr(torch.Tensor, name) and is_method(node, name)) diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py index d5267a26179..d1f5dba0373 100644 --- a/vllm/model_executor/models/transformers/moe.py +++ b/vllm/model_executor/models/transformers/moe.py @@ -26,9 +26,11 @@ import torch.nn as nn from vllm.config.utils import getattr_iter from vllm.distributed import get_dp_group, get_ep_group from vllm.forward_context import ForwardContext, get_forward_context +from vllm.logger import init_logger from vllm.model_executor.custom_op import PluggableLayer from vllm.model_executor.layers.fused_moe import FusedMoE, MoERunner, RoutedExperts from vllm.model_executor.models.interfaces import MixtureOfExperts +from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser from vllm.model_executor.models.utils import maybe_prefix from vllm.utils.torch_utils import direct_register_custom_op @@ -37,6 +39,8 @@ from .utils import log_replacement if TYPE_CHECKING: from vllm.config import VllmConfig +logger = init_logger(__name__) + @dataclass class TransformersMoEState: @@ -142,14 +146,15 @@ class MoEMixin(MixtureOfExperts): self.num_physical_experts = num_physical_experts self.num_local_physical_experts = num_local_physical_experts self.num_redundant_experts = num_physical_experts - self.num_logical_experts - for mlp in self.mlp_layers: - mlp.n_local_physical_experts = num_local_physical_experts - mlp.n_physical_experts = num_physical_experts - mlp.n_redundant_experts = self.num_redundant_experts - mlp.experts.update_expert_map() + for moe_block in self.mlp_layers: + moe_block.n_local_physical_experts = num_local_physical_experts + moe_block.n_physical_experts = num_physical_experts + moe_block.n_redundant_experts = self.num_redundant_experts + moe_block.experts.update_expert_map() def recursive_replace(self): """Initialize the MoE layers.""" + experts_name = "experts" text_config = self.text_config # Positional arguments @@ -217,10 +222,13 @@ class MoEMixin(MixtureOfExperts): # down_proj = (num_experts, intermediate_size, hidden_size) params = list(child_module.parameters()) is_3d = len(params) > 0 and all(p.ndim == 3 for p in params) - if child_name == "experts" and (is_modulelist or is_3d): + if child_name == experts_name and (is_modulelist or is_3d): # Alias for readability - mlp = module + moe_block = module experts = child_module + # Class of the fused block (parent of gate/experts/shared) + moe_block_cls = type(moe_block).__name__ + experts_cls = type(experts).__name__ # Do the experts have biases has_bias = False for experts_param_name, _ in experts.named_parameters(): @@ -230,64 +238,93 @@ class MoEMixin(MixtureOfExperts): # If the config does not specify num_shared_experts, but # the model has shared experts, we assume there is one. if self.num_shared_experts == 0: - for mlp_param_name, _ in mlp.named_parameters(): - if "shared_expert" in mlp_param_name: + for moe_block_param_name, _ in moe_block.named_parameters(): + if "shared_expert" in moe_block_param_name: self.num_shared_experts = 1 break - # Replace experts module with FusedMoE - moe_state = TransformersMoEState() - - def custom_routing_function( - hidden_states: torch.Tensor, - gating_output: torch.Tensor, - topk: int, - renormalize: bool, - moe_state: TransformersMoEState, - ): - """Return `topk_weights` from `gating_output` and the - `topk_ids` we stored in the layer earlier.""" - topk_weights = gating_output - topk_ids = moe_state.topk_ids - assert topk_ids is not None - # Handle all gather in expert parallel - if topk_ids.size(0) != hidden_states.size(0): - dp_metadata = get_forward_context().dp_metadata - sizes = dp_metadata.get_chunk_sizes_across_dp_rank() - is_sp = moe_state.is_sequence_parallel - dist_group = get_ep_group() if is_sp else get_dp_group() - assert sizes[dist_group.rank_in_group] == topk_ids.shape[0] - (topk_ids,) = dist_group.all_gatherv([topk_ids], 0, sizes) - return topk_weights, topk_ids - - fused_experts = FusedMoE( + kwargs: dict[str, Any] = dict( num_experts=num_experts, top_k=top_k, hidden_size=hidden_size, intermediate_size=intermediate_size, renormalize=renormalize, - # Hard coded because topk happens in Transformers use_grouped_topk=False, - num_expert_group=num_expert_group, - topk_group=topk_group, quant_config=self.quant_config, prefix=qual_name, activation=activation, enable_eplb=enable_eplb, num_redundant_experts=num_redundant_experts, has_bias=has_bias, - custom_routing_function=partial( - custom_routing_function, - moe_state=moe_state, - ), - runner_cls=TransformersMoERunner, routed_experts_cls=TransformersRoutedExperts, - runner_args={"moe_state": moe_state}, ) - mlp.experts = fused_experts + fuser = MoEBlockFuser.match(moe_block, experts_name) + if self.num_expert_groups <= 1 and fuser is not None: + # MoE block forward is fully replaced. + # gate/router and shared expert (if any) runs in FusedMoE. + kwargs |= dict( + scoring_func=fuser.scoring_func, + is_sequence_parallel=( + self.parallel_config.use_sequence_parallel_moe + ), + gate=fuser.gate(moe_block, prefix), + shared_experts=fuser.shared_experts(moe_block, prefix), + ) + fuser.rewrite_forward(moe_block) + routed = "gate + experts" + if fuser.shared_name: + routed += " + shared experts" + logger.info_once( + "Fused: %s (%s) -> FusedMoE (internal routing)", + routed, + moe_block_cls, + ) + else: + # MoE block forward is unmodified. + # gate/router and shared expert (if any) runs in Transformers. + # We then smuggle the topk_ids in using a custom op. + moe_state = TransformersMoEState() + + def custom_routing_function( + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + moe_state: TransformersMoEState, + ): + """Return `topk_weights` from `gating_output` and the + `topk_ids` we stored in the layer earlier.""" + topk_weights = gating_output + topk_ids = moe_state.topk_ids + assert topk_ids is not None + # Handle all gather in expert parallel + if topk_ids.size(0) != hidden_states.size(0): + dp_metadata = get_forward_context().dp_metadata + sizes = dp_metadata.get_chunk_sizes_across_dp_rank() + is_sp = moe_state.is_sequence_parallel + group = get_ep_group() if is_sp else get_dp_group() + assert sizes[group.rank_in_group] == topk_ids.shape[0] + (topk_ids,) = group.all_gatherv([topk_ids], 0, sizes) + return topk_weights, topk_ids + + kwargs |= dict( + num_expert_group=num_expert_group, + topk_group=topk_group, + custom_routing_function=partial( + custom_routing_function, moe_state=moe_state + ), + runner_cls=TransformersMoERunner, + runner_args={"moe_state": moe_state}, + ) + logger.info_once( + "Fused: experts (%s) -> FusedMoE (external routing)", + experts_cls, + ) + fused_experts = FusedMoE(**kwargs) + moe_block.experts = fused_experts log_replacement(qual_name, experts, fused_experts) # Update MixtureOfExperts mixin state - self.mlp_layers.append(mlp) + self.mlp_layers.append(moe_block) self.moe_layers.append(fused_experts) else: _recursive_replace(child_module, prefix=qual_name) diff --git a/vllm/model_executor/models/transformers/utils.py b/vllm/model_executor/models/transformers/utils.py index 0a4ca94c5e9..4d9b01ce393 100644 --- a/vllm/model_executor/models/transformers/utils.py +++ b/vllm/model_executor/models/transformers/utils.py @@ -23,10 +23,8 @@ from typing import TYPE_CHECKING, Literal import torch from torch import nn -from vllm.config.utils import getattr_iter from vllm.logger import init_logger from vllm.model_executor.layers.conv import Conv2dLayer, Conv3dLayer -from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, ReplicatedLinear, @@ -183,45 +181,6 @@ def replace_conv_class(conv: TorchConv) -> VllmConv | TorchConv: ) -def replace_rms_norm_class(rms_norm: nn.Module, hidden_size: int) -> RMSNorm: - """Replace a Transformers RMSNorm with vLLM's RMSNorm. - - This method assumes: - - Weight is stored as `weight`. - - Epsilon is stored as `eps` or `variance_epsilon`. - - `with_scale` indicates whether the layer has a weight (Gemma3n only). - - `var_hidden_size` is only ever used for Intern vision encoder in vLLM - and Transformers doesn't appear to have the same concept. - """ - eps = getattr_iter(rms_norm, ("eps", "variance_epsilon"), 1e-6) - kwargs = {"hidden_size": hidden_size, "eps": eps} - # Update hidden size if weight is available - weight_meta = getattr(rms_norm, "weight", None) - if weight_meta is not None: - kwargs["hidden_size"] = weight_meta.size(0) - # Check if weight is all zeros, which indicates GemmaRMSNorm - # We must create a new instance because rms_norm is on meta - try: - with torch.device("cpu"): - weight_test = getattr(rms_norm.__class__(1), "weight", None) - except Exception: - logger.warning( - "Failed to determine if RMSNorm weight is centered on zero or one. " - "Defaulting to one." - ) - weight_test = None - if weight_test is not None and torch.all(weight_test == 0): - return GemmaRMSNorm(**kwargs) - # Otherwise assume it's a regular RMSNorm - kwargs["has_weight"] = getattr(rms_norm, "with_scale", True) - if weight_meta is not None: - kwargs["dtype"] = weight_meta.dtype - else: - # No weight, fall back to weightless RMSNorm - kwargs["has_weight"] = False - return RMSNorm(**kwargs) - - def recursive_replace_linear( model: nn.Module, quant_config: "QuantizationConfig | None", diff --git a/vllm/model_executor/offloader/prefetch.py b/vllm/model_executor/offloader/prefetch.py index 2c860632650..466d8c13ce7 100644 --- a/vllm/model_executor/offloader/prefetch.py +++ b/vllm/model_executor/offloader/prefetch.py @@ -388,7 +388,7 @@ class _ModuleOffloader: # Event to signal when H2D copy to static buffer is complete. # Used for per-layer synchronization (both eager and capture modes). - self._copy_done_event = torch.Event() + self._copy_done_event = torch.cuda.Event() # Track whether _copy_done_event is valid for eager-mode wait_event. # False when: (1) never recorded, or (2) last recorded during a @@ -518,7 +518,7 @@ class _ModuleOffloader: # Fork: record event on compute stream, copy_stream waits on it # This joins copy_stream to any active CUDA graph capture - fork_event = torch.Event() + fork_event = torch.cuda.Event() torch.cuda.current_stream().record_event(fork_event) self.copy_stream.wait_event(fork_event) diff --git a/vllm/model_executor/warmup/deep_gemm_warmup.py b/vllm/model_executor/warmup/deep_gemm_warmup.py index d41604fc7a6..cfff491ab2b 100644 --- a/vllm/model_executor/warmup/deep_gemm_warmup.py +++ b/vllm/model_executor/warmup/deep_gemm_warmup.py @@ -11,6 +11,9 @@ from tqdm import tqdm import vllm.envs as envs from vllm.distributed.parallel_state import get_dp_group, is_global_first_rank +from vllm.model_executor.kernels.linear.scaled_mm.deep_gemm import ( + DeepGemmFp8BlockScaledMMKernel, +) from vllm.model_executor.layers.fused_moe import MoERunner from vllm.model_executor.layers.fused_moe.deep_gemm_utils import ( compute_aligned_M_and_alignment, @@ -147,6 +150,12 @@ def _fp8_linear_may_use_deep_gemm(module: torch.nn.Module) -> bool: ): return False + if not isinstance( + getattr(module.quant_method, "fp8_linear", None), + DeepGemmFp8BlockScaledMMKernel, + ): + return False + w, _, block_sizes = _extract_data_from_linear_base_module(module) return ( block_sizes == get_mk_alignment_for_contiguous_layout() diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index e31a14db663..b7f3c265704 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -148,15 +148,8 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: use_persistent_cache = True - deepep_a2a_backends = { - "deepep_high_throughput", - "deepep_low_latency", - "deepep_v2", - } - if runner.vllm_config.parallel_config.all2all_backend in deepep_a2a_backends: - # DeepEP dispatch/combine can timeout when only rank 0 - # performs autotune and falls behind other ranks. - # Thus we skip persistent cache in this case. + # When distributed, tune on every rank so the collectives stay synchronized. + if get_world_group().world_size > 1: use_persistent_cache = False if not use_persistent_cache: diff --git a/vllm/models/deepseek_v32/nvidia/mtp.py b/vllm/models/deepseek_v32/nvidia/mtp.py index 118f27459bb..baeb8b25e08 100644 --- a/vllm/models/deepseek_v32/nvidia/mtp.py +++ b/vllm/models/deepseek_v32/nvidia/mtp.py @@ -93,10 +93,18 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module): # main model fuses that all-reduce into the next norm, but here the # recycle hidden is consumed directly, so reduce it now. hidden_states = tensor_model_parallel_all_reduce(hidden_states) - # Return the pre-final-norm recycle hidden (re-fed as the next spec - # step's previous_hidden_states); shared_head norm is applied in - # compute_logits. Matches the V2-runner / deepseek_v4 MTP contract. - return residual + hidden_states + # Recycle the POST-final-norm hidden into the next draft step. The + # residual-add is fused into the final RMSNorm so it is computed + # exactly once, and the result is returned for both tuple positions: + # the draft-logits hidden (compute_logits applies the LM head only) and + # the recycled previous_hidden_states. Recycling the pre-final-norm + # hidden mismatches the draft model's hnorm and lowers MTP acceptance; + # post-norm recycle matches deepseek_mtp.py (PR #45895). The tuple form + # is understood by both the V2 speculator (isinstance-tuple check) and + # the legacy proposer (model_returns_tuple is True for the + # DeepSeekMTPModel architecture). + hidden_states, _ = self.shared_head.norm(hidden_states, residual) + return hidden_states, hidden_states class DeepseekV32MultiTokenPredictor(nn.Module): @@ -168,9 +176,10 @@ class DeepseekV32MultiTokenPredictor(nn.Module): ) -> torch.Tensor: current_step_idx = spec_step_idx % self.num_mtp_layers mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)] - return self.logits_processor( - mtp_layer.shared_head.head, mtp_layer.shared_head(hidden_states) - ) + # hidden_states is already post-final-norm (produced in the layer + # forward and recycled as-is); apply the LM head only, without a + # second RMSNorm. + return self.logits_processor(mtp_layer.shared_head.head, hidden_states) class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts): diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 519f5f9a144..906d594ac53 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -56,7 +56,11 @@ from vllm.v1.attention.backends.mla.indexer import ( get_max_prefill_buffer_size, ) from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache -from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec +from vllm.v1.kv_cache_interface import ( + KVCacheSpec, + MLAAttentionSpec, + get_kv_quant_mode, +) logger = init_logger(__name__) @@ -272,7 +276,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events; # [1] doubles as post-GEMM event1. Reuse is safe: GEMM fully joins # before post-GEMM starts. - self.ln_events = [torch.Event() for _ in range(4)] + self.ln_events = [torch.cuda.Event() for _ in range(4)] assert cache_config is not None, "DeepseekV4 attention requires cache_config" # ---- Attention / KV-cache setup ---- @@ -616,6 +620,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): cache_dtype_str=self.kv_cache_dtype, alignment=576 if uses_fp8_ds_mla_layout else None, model_version="deepseek_v4", + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), ) @@ -760,7 +765,10 @@ class DeepseekV4Indexer(nn.Module): # None on ROCm — maybe_execute_in_parallel falls back to sequential. self.aux_stream = aux_stream - self.ln_events: list[torch.Event] = [torch.Event(), torch.Event()] + self.ln_events: list[torch.cuda.Event] = [ + torch.cuda.Event(), + torch.cuda.Event(), + ] def forward( self, diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py index 1efa987fe7b..48cd2340387 100644 --- a/vllm/models/deepseek_v4/compressor.py +++ b/vllm/models/deepseek_v4/compressor.py @@ -104,12 +104,9 @@ class CompressorMetadataBuilder(AttentionMetadataBuilder): common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, ) -> CompressorMetadata: - query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu - num_reqs = common_attn_metadata.num_reqs - query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] - x = torch.repeat_interleave(torch.arange(num_reqs), query_lens).pin_memory() - token_to_req_indices = self.token_to_req_indices[: x.shape[0]] - token_to_req_indices.copy_(x, non_blocking=True) + token_to_req_indices = common_attn_metadata.token_to_req_indices( + self.token_to_req_indices + ) return CompressorMetadata( block_table=common_attn_metadata.block_table_tensor.clamp_(min=0), slot_mapping=common_attn_metadata.slot_mapping, diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index be4a87b323f..e4e258372a2 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -269,6 +269,8 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): # load_dspark_model always aliases the target's. has_own_embed_tokens = False has_own_lm_head = False + # Full-vocab draft: draft ids are target ids, no remapping needed. + draft_id_to_target_id = None def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: super().__init__() diff --git a/vllm/models/deepseek_v4/sparse_mla.py b/vllm/models/deepseek_v4/sparse_mla.py index 136a96a45da..1aaf3f1a141 100644 --- a/vllm/models/deepseek_v4/sparse_mla.py +++ b/vllm/models/deepseek_v4/sparse_mla.py @@ -5,7 +5,6 @@ from dataclasses import dataclass from typing import Any, ClassVar -import numpy as np import torch from vllm.config import VllmConfig @@ -13,7 +12,6 @@ from vllm.config.cache import CacheDType from vllm.platforms.interface import DeviceCapability from vllm.triton_utils import tl, triton from vllm.utils.math_utils import cdiv -from vllm.utils.torch_utils import np_to_pinned_tensor from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -203,18 +201,7 @@ class DeepseekV4FlashMLAMetadataBuilder( fast_build: bool = False, ) -> DeepseekV4FlashMLAMetadata: cm = common_attn_metadata - num_tokens = cm.num_actual_tokens - starts = np.asarray(cm.query_start_loc_cpu, dtype=np.int32) - seg_lengths = np.diff(starts) - req_id_per_token = np.repeat( - np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths - ) - # Zero-fill for cudagraphs - self.req_id_per_token_buffer.fill_(0) - self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_( - np_to_pinned_tensor(req_id_per_token), non_blocking=True - ) - req_id_per_token = self.req_id_per_token_buffer[:num_tokens] + req_id_per_token = cm.token_to_req_indices(self.req_id_per_token_buffer) slot_mapping = cm.slot_mapping if self.compress_ratio > 1: diff --git a/vllm/models/deepseek_v4/xpu/xpu_sparse.py b/vllm/models/deepseek_v4/xpu/xpu_sparse.py index 74d27d7bc41..77cc35cf492 100644 --- a/vllm/models/deepseek_v4/xpu/xpu_sparse.py +++ b/vllm/models/deepseek_v4/xpu/xpu_sparse.py @@ -44,6 +44,17 @@ class DeepseekV4XPUAttention(DeepseekV4Attention): backend_cls = DeepseekV4XPUSparseBackend use_flashmla_fp8_layout = True + def __init__(self, *args, **kwargs) -> None: + # torch.cuda.Event() raises RuntimeError on XPU ("dummy base class"). + # The Base and DeepseekV4Indexer both create cuda Events in __init__, so + # we temporarily redirect torch.cuda.Event → torch.xpu.Event. + _orig_event = torch.cuda.Event + torch.cuda.Event = torch.xpu.Event # type: ignore[misc] + try: + super().__init__(*args, **kwargs) + finally: + torch.cuda.Event = _orig_event # type: ignore[misc] + def _fused_qnorm_rope_kv_insert(self, q, kv, positions, attn_metadata): from typing import cast diff --git a/vllm/multimodal/media/connector.py b/vllm/multimodal/media/connector.py index 582b6fde565..a440e69ef1c 100644 --- a/vllm/multimodal/media/connector.py +++ b/vllm/multimodal/media/connector.py @@ -13,14 +13,17 @@ from pathlib import Path from typing import Any, TypeVar from urllib.request import url2pathname +import aiohttp import numpy as np import numpy.typing as npt +import requests import torch from PIL import Image, UnidentifiedImageError from urllib3.util import Url, parse_url import vllm.envs as envs from vllm.connections import HTTPConnection, global_http_connection +from vllm.exceptions import VLLMUnprocessableEntityError from vllm.logger import init_logger from vllm.multimodal.video import get_video_loader_backend_for_processor from vllm.utils.registry import ExtensionManager @@ -48,6 +51,65 @@ MODALITY_IO_MAP: dict[str, type[MediaIO]] = { } +def _wrap_media_fetch_error( + url: str, exc: Exception +) -> VLLMUnprocessableEntityError | Exception: + """Convert media fetch exceptions to VLLMUnprocessableEntityError. + + This handles HTTP errors that indicate the media resource is invalid + (4xx responses except 408/429, malformed URLs) and converts them to a + 422 Unprocessable Entity error instead of 500. + + Transient errors (5xx, 408, 429, DNS failures, connection errors, + timeouts) are returned as-is to allow retry logic to handle them + appropriately. + + Returns: + VLLMUnprocessableEntityError for permanent client errors (4xx except + 408/429, invalid URL) + Original exception for transient errors (5xx, 408, 429, network blips) + or other exceptions + """ + if isinstance(exc, aiohttp.ClientResponseError): + if exc.status in (408, 429): + return exc + if exc.status < 500: + return VLLMUnprocessableEntityError( + f"Failed to fetch media from URL: HTTP {exc.status} error", + parameter="image_url", + value=url, + ) + return exc + + if isinstance(exc, requests.exceptions.HTTPError): + if exc.response is not None: + status_code = exc.response.status_code + if status_code in (408, 429): + return exc + if status_code < 500: + return VLLMUnprocessableEntityError( + f"Failed to fetch media from URL: HTTP {status_code} error", + parameter="image_url", + value=url, + ) + return exc + + if isinstance(exc, requests.exceptions.InvalidURL): + return VLLMUnprocessableEntityError( + "Failed to fetch media from URL: Invalid URL format", + parameter="image_url", + value=url, + ) + + if isinstance(exc, ValueError): + return VLLMUnprocessableEntityError( + "Failed to fetch media from URL: Invalid URL", + parameter="image_url", + value=url, + ) + return exc + + def merge_media_io_kwargs( defaults: dict[str, dict[str, Any]] | None, overrides: dict[str, dict[str, Any]] | None, @@ -303,11 +365,17 @@ class MediaConnector: return media_io.load_bytes(cached) connection = self.connection - data = connection.get_bytes( - url_spec.url, - timeout=fetch_timeout, - allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS, - ) + try: + data = connection.get_bytes( + url_spec.url, + timeout=fetch_timeout, + allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS, + ) + except Exception as e: + wrapped = _wrap_media_fetch_error(url, e) + if isinstance(wrapped, VLLMUnprocessableEntityError): + raise wrapped from e + raise self._put_cached_bytes(url, data) return media_io.load_bytes(data) @@ -348,11 +416,17 @@ class MediaConnector: return await future connection = self.connection - data = await connection.async_get_bytes( - url_spec.url, - timeout=fetch_timeout, - allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS, - ) + try: + data = await connection.async_get_bytes( + url_spec.url, + timeout=fetch_timeout, + allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS, + ) + except Exception as e: + wrapped = _wrap_media_fetch_error(url, e) + if isinstance(wrapped, VLLMUnprocessableEntityError): + raise wrapped from e + raise await loop.run_in_executor( global_thread_pool, self._put_cached_bytes, url, data diff --git a/vllm/multimodal/media/video.py b/vllm/multimodal/media/video.py index 404f5a0e7cf..45ea4c2fdf4 100644 --- a/vllm/multimodal/media/video.py +++ b/vllm/multimodal/media/video.py @@ -10,11 +10,14 @@ import pybase64 from PIL import Image from vllm import envs +from vllm.logger import init_logger from ..video import VIDEO_LOADER_REGISTRY from .base import MediaIO from .image import ImageMediaIO +logger = init_logger(__name__) + class VideoMediaIO(MediaIO[tuple[npt.NDArray, dict[str, Any]]]): """Configuration values can be user-provided either by --media-io-kwargs or @@ -28,6 +31,24 @@ class VideoMediaIO(MediaIO[tuple[npt.NDArray, dict[str, Any]]]): default_kwargs: dict[str, Any] | None, runtime_kwargs: dict[str, Any] | None, ) -> dict[str, Any]: + if runtime_kwargs: + # Block request-level selection of GPU video backends that + # were not configured (and VRAM-reserved) at startup. + for key in ("video_backend", "backend"): + requested = runtime_kwargs.get(key) + if requested and VIDEO_LOADER_REGISTRY.backend_requires_gpu(requested): + static_val = (default_kwargs or {}).get(key) + if static_val != requested: + logger.warning_once( + "Stripping request-level %s=%r: GPU video " + "backend not configured at startup.", + key, + requested, + ) + runtime_kwargs = { + k: v for k, v in runtime_kwargs.items() if k != key + } + merged = super().merge_kwargs(default_kwargs, runtime_kwargs) # fps and num_frames interact with each other, so if either is # overridden at request time, wipe the other from defaults to diff --git a/vllm/multimodal/video.py b/vllm/multimodal/video.py index 725e33e3f8b..874745e714f 100644 --- a/vllm/multimodal/video.py +++ b/vllm/multimodal/video.py @@ -31,6 +31,13 @@ try: except ImportError: av = PlaceholderModule("av") # type: ignore[assignment] +try: + from torchcodec.decoders import VideoDecoder +except ImportError: + VideoDecoder = PlaceholderModule("torchcodec").placeholder_attr( # type: ignore[assignment] + "decoders.VideoDecoder" + ) + logger = init_logger(__name__) @@ -39,6 +46,7 @@ class VideoLoaderRegistry(ExtensionManager): def __init__(self) -> None: super().__init__() self.processor2backend: dict[str, str] = {} + self._requires_gpu: dict[str, bool] = {} @staticmethod def _normalize_registered_video_processors( @@ -62,11 +70,13 @@ class VideoLoaderRegistry(ExtensionManager): name: str, *, video_processor: str | tuple[str, ...] | None = None, + requires_gpu: bool = False, ): processors = self._normalize_registered_video_processors(video_processor) def wrap(cls_to_register): self.name2class[name] = cls_to_register + self._requires_gpu[name] = requires_gpu for processor_name in processors: self.processor2backend[processor_name] = name return cls_to_register @@ -82,6 +92,9 @@ class VideoLoaderRegistry(ExtensionManager): return self.processor2backend.get(video_processor) + def backend_requires_gpu(self, name: str) -> bool: + return self._requires_gpu.get(name, False) + def get_video_loader_backend_for_processor( video_processor: str | None, @@ -562,6 +575,53 @@ class PyAVVideoBackendMixin: return np.stack(frames_list), valid_indices +class TorchCodecVideoBackendMixin: + """TorchCodec (FFmpeg-backed, PyTorch-native) codec utilities. + + Builds a :class:`~torchcodec.decoders.VideoDecoder` over the in-memory + bytes and extracts the sampled indices with a single batched + ``get_frames_at`` call, while releasing the GIL during decode. + """ + + @staticmethod + def make_torchcodec_decoder( + data: bytes, + *, + num_ffmpeg_threads: int = 0, + seek_mode: Literal["exact", "approximate"] = "exact", + ) -> "VideoDecoder": + # NHWC matches the (num_frames, H, W, 3) uint8 RGB layout the rest + # of the pipeline expects, avoiding a transpose. + return VideoDecoder( + data, + dimension_order="NHWC", + num_ffmpeg_threads=num_ffmpeg_threads, + seek_mode=seek_mode, + ) + + @staticmethod + def get_torchcodec_metadata(decoder: "VideoDecoder") -> VideoSourceMetadata: + md = decoder.metadata + total_frames = md.num_frames or 0 + fps = float(md.average_fps) if md.average_fps else 0.0 + duration = float(md.duration_seconds) if md.duration_seconds else 0.0 + if total_frames == 0 and duration > 0 and fps > 0: + total_frames = int(duration * fps) + return VideoSourceMetadata(total_frames, fps, duration) + + @staticmethod + def decode_torchcodec_frames( + decoder: "VideoDecoder", + frame_indices: list[int], + ) -> tuple[npt.NDArray, list[int]]: + """Decode the requested indices in one batched, index-exact call.""" + if not frame_indices: + return np.empty((0,), dtype=np.uint8), [] + # Note: torchcodec releases the GIL for the entire call + batch = decoder.get_frames_at(frame_indices) + return batch.data.numpy(), list(frame_indices) + + class PyNvVideoCodecVideoBackendMixin: """PyNvVideoCodec utilities for GPU-backed frame decode.""" @@ -771,14 +831,15 @@ class VideoBackend( VideoLoader, OpenCVVideoBackendMixin, PyAVVideoBackendMixin, + TorchCodecVideoBackendMixin, PyNvVideoCodecVideoBackendMixin, ): """Uniform-sampling video backend. Samples ``num_frames`` uniformly across the video (or one frame every ``1/fps`` seconds, whichever produces fewer frames). The decoding codec - is selected via the ``backend`` kwarg (``"opencv"``, ``"pyav"``, or - ``"pynvvideocodec"``), which can be passed through + is selected via the ``backend`` kwarg (``"opencv"``, ``"pyav"``, + ``"torchcodec"`` or ``"pynvvideocodec"``), which can be passed through ``--media-io-kwargs``. Defaults to ``"opencv"``. """ @@ -824,7 +885,9 @@ class VideoBackend( max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav", "pynvvideocodec"] = "opencv", + backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec"] = "opencv", + num_ffmpeg_threads: int = 0, + seek_mode: Literal["exact", "approximate"] = "exact", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: """Load sampled frames from raw video bytes. @@ -837,8 +900,20 @@ class VideoBackend( dynamic subclass; ignored here. frame_recovery: Enable forward-scan recovery for failed frames. Only honored by the OpenCV codec. - backend: Decoding codec — ``"opencv"``, ``"pyav"``, or - ``"pynvvideocodec"``. + backend: Decoding codec — ``"opencv"``, ``"pyav"``, + ``"torchcodec"`` or ``"pynvvideocodec"``. + num_ffmpeg_threads: Number of FFmpeg decoding threads, only used by + TorchCodec: ``0`` (default) relies on the FFmpeg default value + which is ``min(cpu_count + 1, 16)``. + OpenCV will always use ``min(cpu_count, 16)`` while pyav will + always use ``min(cpu_count, (height + 15) / 16)``. + seek_mode: Seek mode for the TorchCodec decoder, only used by + TorchCodec: ``"exact"`` (default) guarantees frame-accurate + sampling by scanning the file on creation, while + ``"approximate"`` skips that scan for faster decoder creation + at the cost of relying on the file's metadata. See + https://meta-pytorch.org/torchcodec/stable/generated_examples/decoding/approximate_mode.html + for details. Returns: Tuple of ``(frames_array, metadata_dict)``. @@ -877,6 +952,24 @@ class VideoBackend( frames, valid = cls.decode_frames( container, frame_idx, source.original_fps, source.duration ) + elif backend == "torchcodec": + assert not frame_recovery, ( + "frame_recovery is only available for `opencv` backend" + ) + decoder = cls.make_torchcodec_decoder( + data, + num_ffmpeg_threads=num_ffmpeg_threads, + seek_mode=seek_mode, + ) + _check_frame_pixel_limit( + decoder.metadata.width or 0, + decoder.metadata.height or 0, + ) + source = cls._prepare_source(cls.get_torchcodec_metadata(decoder)) + frame_idx = cls.compute_frames_index_to_sample( + source=source, target=target, **kwargs + ) + frames, valid = cls.decode_torchcodec_frames(decoder, frame_idx) elif backend == PYNVVIDEOCODEC_VIDEO_BACKEND: if frame_recovery: raise ValueError( @@ -891,7 +984,8 @@ class VideoBackend( else: raise ValueError( f"Unknown video codec backend {backend!r}; " - "valid options: 'opencv', 'pyav', 'pynvvideocodec'." + "valid options: 'opencv', 'pyav', 'torchcodec', " + "'pynvvideocodec'." ) if len(valid) < len(frame_idx): @@ -909,7 +1003,7 @@ class VideoBackend( ) -@VIDEO_LOADER_REGISTRY.register(PYNVVIDEOCODEC_VIDEO_BACKEND) +@VIDEO_LOADER_REGISTRY.register(PYNVVIDEOCODEC_VIDEO_BACKEND, requires_gpu=True) class PyNvVideoCodecVideoBackend(VideoBackend): """Hardware-accelerated video backend using PyNvVideoCodec. @@ -978,7 +1072,7 @@ class Qwen3VLVideoBackend(VideoBackend): max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav", "pynvvideocodec"] = "opencv", + backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec"] = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: return super().load_bytes( @@ -1057,7 +1151,7 @@ class Qwen2VLVideoBackend(VideoBackend): max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav", "pynvvideocodec"] = "opencv", + backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec"] = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: return super().load_bytes( @@ -1149,7 +1243,7 @@ class DynamicVideoBackend(VideoBackend): max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav", "pynvvideocodec"] = "opencv", + backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec"] = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: return super().load_bytes( @@ -1274,7 +1368,7 @@ class GLM46VVideoBackend(VideoBackend): max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav", "pynvvideocodec"] = "opencv", + backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec"] = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: return super().load_bytes( @@ -1372,7 +1466,7 @@ class GLMGAVideoBackend(VideoBackend): max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav", "pynvvideocodec"] = "opencv", + backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec"] = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: frames, metadata = super().load_bytes( @@ -1695,7 +1789,7 @@ class NemotronVLVideoBackend(VideoBackend): max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav", "pynvvideocodec"] = "opencv", + backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec"] = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: frames, metadata = super().load_bytes( diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index 867833c9d7e..90a5a6b9219 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -148,6 +148,15 @@ class XPUPlatform(Platform): if selected_backend == AttentionBackendEnum.TRITON_ATTN: logger.info_once("Using Triton backend.") return AttentionBackendEnum.TRITON_ATTN.get_path() + elif attn_selector_config.use_mm_prefix: + # Flash Attention on XPU has no FA4 kernel, so it cannot apply the + # multimodal prefix-LM bidirectional mask. Fall back to Triton + # Attention, which supports mm_prefix. + logger.warning_once( + "Flash Attention on XPU does not support multimodal prefix-LM " + "attention. Falling back to Triton Attention backend." + ) + return AttentionBackendEnum.TRITON_ATTN.get_path() elif dtype == torch.float32: logger.warning_once( "Flash Attention on XPU does not support float32 dtype. " diff --git a/vllm/tool_parsers/granite_tool_parser.py b/vllm/tool_parsers/granite_tool_parser.py index d586db32670..174e2884277 100644 --- a/vllm/tool_parsers/granite_tool_parser.py +++ b/vllm/tool_parsers/granite_tool_parser.py @@ -154,9 +154,11 @@ class GraniteToolParser(ToolParser): current_tool_call: dict = tool_call_arr[self.current_tool_id] delta = None - # case: we are starting a new tool in the array - # -> array has > 0 length AND length has moved past cursor - if len(tool_call_arr) > self.current_tool_id + 1: + # Only advance once the current tool name is streamed; granite + # emits arguments before name, so advancing early would drop it. + if len(tool_call_arr) > self.current_tool_id + 1 and ( + self.current_tool_id < 0 or self.current_tool_name_sent + ): # if we're moving on to a new call, first make sure we # haven't missed anything in the previous one that was # auto-generated due to JSON completions, but wasn't @@ -184,7 +186,7 @@ class GraniteToolParser(ToolParser): ) # re-set stuff pertaining to progress in the current tool - self.current_tool_id = len(tool_call_arr) - 1 + self.current_tool_id += 1 self.current_tool_name_sent = False self.streamed_args_for_tool.append("") logger.debug("starting on new tool %d", self.current_tool_id) diff --git a/vllm/tool_parsers/mistral_tool_parser.py b/vllm/tool_parsers/mistral_tool_parser.py index 1d605557b1f..026098a8735 100644 --- a/vllm/tool_parsers/mistral_tool_parser.py +++ b/vllm/tool_parsers/mistral_tool_parser.py @@ -533,6 +533,7 @@ class MistralToolParser(ToolParser): if prefix == "item" and event == "start_map": self.streaming_state = StreamingState.WAITING_FOR_TOOL_KEY + self.starting_new_tool = True if prefix == "item" and event == "map_key" and value == "name": self.streaming_state = StreamingState.PARSING_NAME if prefix == "item.name" and event == "string": @@ -640,18 +641,10 @@ class MistralToolParser(ToolParser): # Given the parsed text and the possible streaming state change, # let's add to the tool delta - if ( - (streaming_state_before_parse != self.streaming_state) - and streaming_state_before_parse - in [StreamingState.WAITING_FOR_TOOL_START, StreamingState.TOOL_COMPLETE] - and self.streaming_state - not in [ - StreamingState.ALL_TOOLS_COMPLETE, - StreamingState.TOOL_COMPLETE, - StreamingState.WAITING_FOR_TOOL_START, - ] - ): - # starting a new tool call + # start_map is the authoritative new-tool signal and survives + # batched deltas, unlike comparing pre/post streaming states + if self.starting_new_tool: + self.starting_new_tool = False if current_tool_call_modified: if self.current_tool_mistral_id is not None: current_tool_call.id = self.current_tool_mistral_id diff --git a/vllm/tool_parsers/utils.py b/vllm/tool_parsers/utils.py index a31420cf1cd..95769bafd7f 100644 --- a/vllm/tool_parsers/utils.py +++ b/vllm/tool_parsers/utils.py @@ -5,12 +5,14 @@ import ast import json import math import warnings +from dataclasses import dataclass from json import JSONDecodeError, JSONDecoder from typing import Any, TypeAlias import partial_json_parser from openai.types.responses import ( FunctionTool, + NamespaceTool, ToolChoiceFunction, ) from openai.types.responses.tool import Tool as ResponsesTool @@ -166,6 +168,91 @@ def consume_space(i: int, s: str) -> int: return i +_NAMESPACE_TOOL_SEPARATOR = "__" + + +@dataclass(frozen=True) +class ResponsesToolCallName: + name: str + namespace: str | None = None + + +def flat_namespace_tool_name(namespace: str, name: str) -> str: + return f"{namespace}{_NAMESPACE_TOOL_SEPARATOR}{name}" + + +def iter_response_function_tool_info( + tool: ResponsesTool, +) -> list[tuple[str, dict[str, Any] | None]]: + if isinstance(tool, FunctionTool): + return [(tool.name, tool.parameters)] + if not isinstance(tool, NamespaceTool): + return [] + + namespace = tool.name + return [ + ( + flat_namespace_tool_name(namespace, namespaced_tool.name), + namespaced_tool.parameters, + ) + for namespaced_tool in tool.tools + if namespaced_tool.type == "function" + ] + + +def iter_response_function_tool_dicts( + tools: list[ResponsesTool], +) -> list[dict[str, Any]]: + function_tools: list[dict[str, Any]] = [] + for tool in tools: + if isinstance(tool, NamespaceTool): + namespace = tool.name + for namespaced_tool in tool.tools: + if namespaced_tool.type != "function": + continue + tool_dict = namespaced_tool.model_dump() + tool_dict["name"] = flat_namespace_tool_name( + namespace, namespaced_tool.name + ) + function_tools.append(tool_dict) + else: + function_tools.append(tool.model_dump()) + return function_tools + + +def build_responses_tool_call_name_map( + tools: list[ResponsesTool] | None, +) -> dict[str, ResponsesToolCallName]: + if not tools: + return {} + + name_map: dict[str, ResponsesToolCallName] = {} + for tool in tools: + if not isinstance(tool, NamespaceTool): + continue + namespace = tool.name + for namespaced_tool in tool.tools: + if namespaced_tool.type != "function": + continue + flat_name = flat_namespace_tool_name(namespace, namespaced_tool.name) + name_map[flat_name] = ResponsesToolCallName( + name=namespaced_tool.name, + namespace=namespace, + ) + return name_map + + +def resolve_responses_tool_call_name( + name: str, + tools: list[ResponsesTool] | None = None, + tool_call_name_map: dict[str, ResponsesToolCallName] | None = None, +) -> ResponsesToolCallName: + name_map = tool_call_name_map + if name_map is None: + name_map = build_responses_tool_call_name_map(tools) + return name_map.get(name, ResponsesToolCallName(name=name)) + + def _is_function_tool(tool: Tool) -> bool: return isinstance(tool, (FunctionTool, ChatCompletionToolsParam)) @@ -189,6 +276,11 @@ def find_tool_properties( if not tools: return {} for tool in tools: + if isinstance(tool, (FunctionTool, NamespaceTool)): + for name, params in iter_response_function_tool_info(tool): + if name == tool_name: + return (params or {}).get("properties", {}) + continue if not _is_function_tool(tool): continue name, params = _extract_tool_info(tool) @@ -205,6 +297,11 @@ def find_tool_name( if not tools: return False for tool in tools: + if isinstance(tool, (FunctionTool, NamespaceTool)): + for name, _ in iter_response_function_tool_info(tool): + if name == tool_name: + return True + continue if not _is_function_tool(tool): continue name, _ = _extract_tool_info(tool) @@ -213,8 +310,9 @@ def find_tool_name( return False -def _get_tool_schema_from_tool(tool: Tool) -> dict: - name, params = _extract_tool_info(tool) +def _get_tool_schema_from_name_and_params( + name: str, params: dict[str, Any] | None +) -> dict: params = params if params else {"type": "object", "properties": {}} return { "properties": { @@ -225,6 +323,11 @@ def _get_tool_schema_from_tool(tool: Tool) -> dict: } +def _get_tool_schema_from_tool(tool: Tool) -> dict: + name, params = _extract_tool_info(tool) + return _get_tool_schema_from_name_and_params(name, params) + + def _get_tool_schema_defs( tools: list[Tool], ) -> dict: @@ -247,13 +350,25 @@ def _get_tool_schema_defs( def _get_json_schema_from_tools( tools: list[Tool], ) -> dict: - fn_tools = [t for t in tools if _is_function_tool(t)] + fn_tool_schemas: list[dict[str, Any]] = [] + fn_tools: list[Tool] = [] + for tool in tools: + if isinstance(tool, (FunctionTool, NamespaceTool)): + fn_tool_schemas.extend( + _get_tool_schema_from_name_and_params(name, params) + for name, params in iter_response_function_tool_info(tool) + ) + if isinstance(tool, FunctionTool): + fn_tools.append(tool) + elif _is_function_tool(tool): + fn_tool_schemas.append(_get_tool_schema_from_tool(tool)) + fn_tools.append(tool) json_schema = { "type": "array", "minItems": 1, "items": { "type": "object", - "anyOf": [_get_tool_schema_from_tool(tool) for tool in fn_tools], + "anyOf": fn_tool_schemas, }, } json_schema_defs = _get_tool_schema_defs(fn_tools) @@ -274,23 +389,30 @@ def get_json_schema_from_tools( tool_choice, ToolChoiceFunction ): tool_name = tool_choice.name - tool_map = {tool.name: tool for tool in tools if isinstance(tool, FunctionTool)} - if tool_name not in tool_map: + responses_tool_map: dict[str, dict[str, Any] | None] = {} + for tool in tools: + if not isinstance(tool, (FunctionTool, NamespaceTool)): + continue + for name, params in iter_response_function_tool_info(tool): + responses_tool_map[name] = params + if "__" in name: + responses_tool_map.setdefault(name.rsplit("__", 1)[1], params) + if tool_name not in responses_tool_map: raise ValueError(f"Tool '{tool_name}' has not been passed in `tools`.") - return tool_map[tool_name].parameters + return responses_tool_map[tool_name] # tool_choice: Forced Function (ChatCompletion) if (not isinstance(tool_choice, str)) and isinstance( tool_choice, ChatCompletionNamedToolChoiceParam ): tool_name = tool_choice.function.name - tool_map = { + chat_tool_map: dict[str, ChatCompletionToolsParam] = { tool.function.name: tool for tool in tools if isinstance(tool, ChatCompletionToolsParam) } - if tool_name not in tool_map: + if tool_name not in chat_tool_map: raise ValueError(f"Tool '{tool_name}' has not been passed in `tools`.") - return tool_map[tool_name].function.parameters + return chat_tool_map[tool_name].function.parameters # tool_choice: "required" if tool_choice == "required": return _get_json_schema_from_tools(tools) diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 3b39c911095..7ced083afb5 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -271,6 +271,7 @@ class ModelArchConfigConvertorBase: "pangu_ultra_moe", "pangu_ultra_moe_mtp", "bailing_hybrid", + "bailing_hybrid_mtp", ): # check is deepseek_v4 model if hasattr(self.hf_text_config, "compress_ratios"): @@ -539,6 +540,11 @@ class Qwen3NextMTPModelArchConfigConvertor(ModelArchConfigConvertorBase): return getattr(self.hf_text_config, "num_nextn_predict_layers", 0) +class BailingHybridMTPModelArchConfigConvertor(ModelArchConfigConvertorBase): + def get_num_hidden_layers(self) -> int: + return getattr(self.hf_text_config, "num_nextn_predict_layers", 0) + + class Qwen3_5MTPModelArchConfigConvertor(ModelArchConfigConvertorBase): def get_num_hidden_layers(self) -> int: return getattr(self.hf_text_config, "mtp_num_hidden_layers", 0) @@ -633,6 +639,7 @@ class MossAudioModelArchConfigConvertor(ModelArchConfigConvertorBase): # hf_config.model_type -> convertor class MODEL_ARCH_CONFIG_CONVERTORS = { + "bailing_hybrid_mtp": BailingHybridMTPModelArchConfigConvertor, "cohere_asr": CohereAsrModelArchConfigConvertor, "dbrx": DbrxModelArchConfigConvertor, "deepseek_mtp": DeepSeekMTPModelArchConfigConvertor, diff --git a/vllm/transformers_utils/processor.py b/vllm/transformers_utils/processor.py index fa4c558a739..aa33faed916 100644 --- a/vllm/transformers_utils/processor.py +++ b/vllm/transformers_utils/processor.py @@ -81,6 +81,7 @@ _transformers_v4_compatibility_import() _transformers_v4_compatibility_init() _P = TypeVar("_P", bound=ProcessorMixin, default=ProcessorMixin) +_I = TypeVar("_I", bound=BaseImageProcessor, default=BaseImageProcessor) _V = TypeVar("_V", bound=BaseVideoProcessor, default=BaseVideoProcessor) @@ -440,12 +441,14 @@ def get_image_processor( *args: Any, revision: str | None = None, trust_remote_code: bool = False, + processor_cls_overrides: type[_I] | None = None, **kwargs: Any, ): """Load an image processor for the given model name via HuggingFace.""" try: processor_name = convert_model_repo_to_path(processor_name) - processor = AutoImageProcessor.from_pretrained( + processor_cls = processor_cls_overrides or AutoImageProcessor + processor = processor_cls.from_pretrained( processor_name, *args, revision=revision, diff --git a/vllm/transformers_utils/processors/kimi_k25_vision_fused.py b/vllm/transformers_utils/processors/kimi_k25_vision_fused.py new file mode 100644 index 00000000000..63907898175 --- /dev/null +++ b/vllm/transformers_utils/processors/kimi_k25_vision_fused.py @@ -0,0 +1,352 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Optimized CPU image processor for Kimi-K2.5/K2.6 vision chunks.""" + +import io +import json +import math +from typing import Any + +import numpy as np +import pybase64 as base64 +import torch +from PIL import Image +from transformers.image_processing_utils import BaseImageProcessor, BatchFeature +from transformers.utils import TensorType + +from vllm.utils.import_utils import is_numba_available +from vllm.utils.jit_monitor import numba_workqueue_threading_layer + +if is_numba_available(): + from numba import njit, prange + + @njit(parallel=True, cache=True) + def _write_fused_patches( + frames: np.ndarray, + out: np.ndarray, + out_offset: int, + new_h: int, + new_w: int, + padded_h: int, + padded_w: int, + patch_size: int, + normalize_lut: np.ndarray, + ) -> None: + # frames: [T, new_h, new_w, 3] uint8, without padding. + # out: [total_patches, 3, patch_size, patch_size] float32. + t_size = frames.shape[0] + patch_h = padded_h // patch_size + patch_w = padded_w // patch_size + total = t_size * padded_h * padded_w * 3 + hwc = padded_h * padded_w * 3 + wc = padded_w * 3 + + for linear in prange(total): + t = linear // hwc + rem = linear - t * hwc + y = rem // wc + rem = rem - y * wc + x = rem // 3 + c = rem - x * 3 + + value = frames[t, y, x, c] if y < new_h and x < new_w else 0 + + patch_idx = ( + out_offset + + t * patch_h * patch_w + + (y // patch_size) * patch_w + + (x // patch_size) + ) + out[patch_idx, c, y % patch_size, x % patch_size] = normalize_lut[value, c] + +else: + + def _write_fused_patches(*args: Any, **kwargs: Any) -> None: + raise RuntimeError("numba is required for fused Kimi image preprocessing") + + +def navit_resize_image( + width: int, + height: int, + patch_size: int, + merge_kernel_size: int, + in_patch_limit: int, + patch_limit_on_one_side: int, + fixed_output_tokens: int | None, +) -> dict[str, int]: + s1 = math.sqrt( + in_patch_limit + / (max(1.0, width // patch_size) * max(1.0, height // patch_size)) + ) + s2 = patch_limit_on_one_side * patch_size / width + s3 = patch_limit_on_one_side * patch_size / height + scale = min(1.0, s1, s2, s3) + new_w = min(max(1, int(width * scale)), patch_limit_on_one_side * patch_size) + new_h = min(max(1, int(height * scale)), patch_limit_on_one_side * patch_size) + + factor = merge_kernel_size * patch_size + pad_height = (factor - new_h % factor) % factor + pad_width = (factor - new_w % factor) % factor + + if fixed_output_tokens is not None: + num_tokens = fixed_output_tokens + else: + token_height = (new_h + pad_height) // factor + token_width = (new_w + pad_width) // factor + num_tokens = token_height * token_width + + return { + "num_tokens": num_tokens, + "new_width": new_w, + "new_height": new_h, + "pad_width": pad_width, + "pad_height": pad_height, + "sampled_nframes": 1, + } + + +def navit_resize_video( + width: int, + height: int, + nframes: int, + avg_fps: float, + sample_fps: float, + patch_size: int, + merge_kernel_size: int, + in_patch_limit_each_frame: int, + patch_limit_on_one_side: int, + in_patch_limit_total: int | None, + max_num_frames_each_video: int | None, + fixed_output_tokens_each_frame: int | None, +) -> dict[str, int]: + sample_fps = min(sample_fps, avg_fps) + sampled_nframes = max(round(nframes * sample_fps / avg_fps), 1) + if max_num_frames_each_video is not None: + sampled_nframes = min(sampled_nframes, max_num_frames_each_video) + + if in_patch_limit_total is not None: + in_patch_limit_each_frame = min( + round(in_patch_limit_total / sampled_nframes), + in_patch_limit_each_frame, + ) + + ret = navit_resize_image( + width, + height, + patch_size, + merge_kernel_size, + in_patch_limit_each_frame, + patch_limit_on_one_side, + fixed_output_tokens_each_frame, + ) + ret["sampled_nframes"] = sampled_nframes + return ret + + +def _to_pil(data: Any) -> Image.Image: + if hasattr(data, "media") and hasattr(data, "original_bytes"): + data = data.media + if isinstance(data, Image.Image): + return data if data.mode == "RGB" else data.convert("RGB") + if isinstance(data, str): + if data.startswith("data:"): + raw_base64 = data.split(",", 1)[1] + return Image.open(io.BytesIO(base64.b64decode(raw_base64))).convert("RGB") + return Image.open(data).convert("RGB") + if isinstance(data, bytes): + return Image.open(io.BytesIO(data)).convert("RGB") + raise ValueError(f"Unsupported data type: {type(data)}") + + +def _ensure_media_type(media: dict[str, Any]) -> dict[str, Any]: + if media["type"] == "image": + media["image"] = _to_pil(media["image"]) + return media + if media["type"] == "video_chunk": + media["video_chunk"] = [_to_pil(frame) for frame in media["video_chunk"]] + return media + raise ValueError(f"Unsupported media type: {media['type']}") + + +class KimiK25FusedVisionProcessor(BaseImageProcessor): + model_type = "kimi_k25" + + def __init__(self, media_proc_cfg: dict[str, Any], **kwargs: Any) -> None: + super().__init__(**kwargs) + media_proc_cfg = dict(media_proc_cfg) + merge_kernel_size = media_proc_cfg["merge_kernel_size"] + if isinstance(merge_kernel_size, (list, tuple)): + media_proc_cfg["merge_kernel_size"] = int(merge_kernel_size[0]) + self.media_proc_cfg = media_proc_cfg + self.num_frames_per_chunk = media_proc_cfg["temporal_merge_kernel_size"] + values = np.arange(256, dtype=np.float32)[:, None] + image_mean = np.asarray(media_proc_cfg["image_mean"], dtype=np.float32) + image_std_inv = 1.0 / np.asarray(media_proc_cfg["image_std"], dtype=np.float32) + self.normalize_lut = (values / 255.0 - image_mean[None, :]) * image_std_inv[ + None, : + ] + + def media_tokens_calculator(self, media: dict[str, Any]) -> int: + media = _ensure_media_type(media) + ret = self.get_resize_config(media) + return ret["num_tokens"] + + def get_resize_config(self, media_input: dict[str, Any]) -> dict[str, int]: + if media_input["type"] == "image": + width, height = media_input["image"].size + return navit_resize_image( + width, + height, + self.media_proc_cfg["patch_size"], + self.media_proc_cfg["merge_kernel_size"], + self.media_proc_cfg["in_patch_limit"], + self.media_proc_cfg["patch_limit_on_one_side"], + self.media_proc_cfg["fixed_output_tokens"], + ) + + if media_input["type"] == "video_chunk": + frame = media_input["video_chunk"][0] + width, height = frame.size + num_frames = len(media_input["video_chunk"]) + in_patch_limit_each_frame = self.media_proc_cfg["in_patch_limit_each_frame"] + if in_patch_limit_each_frame is None: + in_patch_limit_each_frame = self.media_proc_cfg["in_patch_limit"] + + return navit_resize_video( + width, + height, + num_frames, + 1.0, + math.inf, + self.media_proc_cfg["patch_size"], + self.media_proc_cfg["merge_kernel_size"], + in_patch_limit_each_frame, + self.media_proc_cfg["patch_limit_on_one_side"], + self.media_proc_cfg["in_patch_limit_video"], + None, + self.media_proc_cfg["fixed_output_tokens"], + ) + + raise ValueError(f"Unsupported type: {media_input['type']}") + + @staticmethod + def resize_image(image: Image.Image, new_width: int, new_height: int) -> np.ndarray: + image = image.resize((new_width, new_height), resample=Image.Resampling.BICUBIC) + return np.asarray(image) + + def preprocess( + self, + medias: list[dict[str, Any]], + return_tensors: str | TensorType | None = None, + ) -> BatchFeature: + if not isinstance(medias, list): + medias = [medias] + if not medias: + return BatchFeature(data={}, tensor_type=return_tensors) + + if njit is None: + raise RuntimeError("numba is required for fused Kimi image preprocessing") + + patch_size = int(self.media_proc_cfg["patch_size"]) + prepared = [] + grid_thws_np = np.empty((len(medias), 3), dtype=np.int64) + total_patches = 0 + + for idx, item in enumerate(medias): + item = _ensure_media_type(item) + resize_config = self.get_resize_config(item) + new_width = resize_config["new_width"] + new_height = resize_config["new_height"] + pad_width = resize_config["pad_width"] + pad_height = resize_config["pad_height"] + padded_width = new_width + pad_width + padded_height = new_height + pad_height + + if item["type"] == "image": + image_np = self.resize_image(item["image"], new_width, new_height) + frames = image_np[np.newaxis, ...] + elif item["type"] == "video_chunk": + frames = np.stack( + [ + self.resize_image(frame, new_width, new_height) + for frame in item["video_chunk"] + ], + axis=0, + ) + else: + raise ValueError(f"Unsupported type: {item['type']}") + + t_size = frames.shape[0] + grid_h = padded_height // patch_size + grid_w = padded_width // patch_size + grid_thws_np[idx, 0] = t_size + grid_thws_np[idx, 1] = grid_h + grid_thws_np[idx, 2] = grid_w + + num_patches = t_size * grid_h * grid_w + prepared.append( + ( + frames, + new_height, + new_width, + padded_height, + padded_width, + num_patches, + ) + ) + total_patches += num_patches + + pixel_values_np = np.empty( + (total_patches, 3, patch_size, patch_size), dtype=np.float32 + ) + out_offset = 0 + with numba_workqueue_threading_layer(): + for ( + frames, + new_height, + new_width, + padded_height, + padded_width, + num_patches, + ) in prepared: + _write_fused_patches( + frames, + pixel_values_np, + out_offset, + new_height, + new_width, + padded_height, + padded_width, + patch_size, + self.normalize_lut, + ) + out_offset += num_patches + + data = { + "pixel_values": torch.from_numpy(pixel_values_np), + "grid_thws": torch.from_numpy(grid_thws_np), + } + return BatchFeature(data=data, tensor_type=return_tensors) + + def __repr__(self): + return f"KimiK25FusedVisionProcessor(media_proc_cfg={self.media_proc_cfg})" + + def to_dict(self) -> dict[str, Any]: + output = super().to_dict() + output["media_proc_cfg"] = self.media_proc_cfg + if "media_processor" in output: + del output["media_processor"] + return output + + @classmethod + def from_dict(cls, config_dict: dict[str, Any], **kwargs): + config = config_dict.copy() + media_proc_cfg = config.pop("media_proc_cfg", {}) + return cls(media_proc_cfg=media_proc_cfg, **config, **kwargs) + + def to_json_string(self): + dictionary = self.to_dict() + for key, value in dictionary.items(): + if hasattr(value, "tolist"): + dictionary[key] = value.tolist() + return json.dumps(dictionary, indent=2, sort_keys=True) + "\n" diff --git a/vllm/utils/import_utils.py b/vllm/utils/import_utils.py index 043798a584b..78dcde6f061 100644 --- a/vllm/utils/import_utils.py +++ b/vllm/utils/import_utils.py @@ -487,6 +487,11 @@ def has_nixl_ep() -> bool: return _has_module("nixl_ep") +def is_numba_available() -> bool: + """Whether the optional `numba` package is available.""" + return _has_module("numba") + + def has_triton_kernels() -> bool: """Whether the optional `triton_kernels` package is available.""" is_available = _has_module("triton_kernels") or _has_module( diff --git a/vllm/utils/jit_monitor.py b/vllm/utils/jit_monitor.py index 4565ffdae06..7ba8ecde653 100644 --- a/vllm/utils/jit_monitor.py +++ b/vllm/utils/jit_monitor.py @@ -20,9 +20,10 @@ Currently monitors: (via ``knobs.runtime.jit_post_compile_hook``) """ +import contextlib import functools import os -from collections.abc import Mapping +from collections.abc import Iterator, Mapping from typing import Literal from vllm.logger import init_logger @@ -287,3 +288,32 @@ def _setup_cutedsl_jit_hook() -> None: cute.compile = _compile_with_monitor _cutedsl_hook_installed = True + + +@contextlib.contextmanager +def numba_workqueue_threading_layer() -> Iterator[None]: + """Force numba's fork-safe `workqueue` threading layer for this block. + + GNU OpenMP (numba's default `omp` threading layer) aborts the process + if a forked child re-enters an OpenMP-active runtime. vLLM forks the + EngineCore subprocess from a process that may already have launched + numba's parallel accelerator, so the first call to any + `@njit(parallel=True)` function must happen under `workqueue` instead. + The threading layer choice is sticky for the life of the process once + launched, so restoring the config on exit does not undo the effect. + """ + import numba + + key = "NUMBA_THREADING_LAYER" + previous_env = os.environ.get(key) + previous_config = numba.config.THREADING_LAYER + os.environ[key] = "workqueue" + numba.config.THREADING_LAYER = "workqueue" + try: + yield + finally: + if previous_env is None: + os.environ.pop(key, None) + else: + os.environ[key] = previous_env + numba.config.THREADING_LAYER = previous_config diff --git a/vllm/utils/multi_stream_utils.py b/vllm/utils/multi_stream_utils.py index fed38ea1a35..2203221c5a1 100644 --- a/vllm/utils/multi_stream_utils.py +++ b/vllm/utils/multi_stream_utils.py @@ -20,8 +20,8 @@ class EventType(Enum): def maybe_execute_in_parallel( fn0: Callable[[], Any], fn1: Callable[[], Any], - event0: torch.Event, - event1: torch.Event, + event0: torch.cuda.Event, + event1: torch.cuda.Event, aux_stream: torch.cuda.Stream | None = None, ) -> tuple[Any, Any]: """Run two functions potentially in parallel on separate CUDA streams. @@ -61,8 +61,8 @@ def maybe_execute_in_parallel( def execute_in_parallel( default_fn: Callable[[], Any], aux_fns: list[Callable[[], Any] | None], - start_event: torch.Event, - done_events: list[torch.Event], + start_event: torch.cuda.Event, + done_events: list[torch.cuda.Event], aux_streams: list[torch.cuda.Stream] | None = None, enable: bool = False, ) -> tuple[Any, list[Any]]: @@ -108,7 +108,7 @@ def execute_in_parallel( ) aux_results = [None] * len(aux_fns) - pending: list[torch.Event] = [] + pending: list[torch.cuda.Event] = [] start_event.record() for i, fn in enumerate(aux_fns): diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 22c6a382287..bfecb3c952e 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -16,6 +16,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8StaticTensorSym, kNvfp4Dynamic, ) +from vllm.utils.torch_utils import np_to_pinned_tensor if TYPE_CHECKING: from vllm.config import VllmConfig @@ -467,6 +468,7 @@ class CommonAttentionMetadata: _num_computed_tokens_cpu: torch.Tensor | None = None _num_computed_tokens_cache: torch.Tensor | None = None + _token_to_req_indices_cache: torch.Tensor | None = None def batch_size(self) -> int: return self.seq_lens.shape[0] @@ -515,6 +517,31 @@ class CommonAttentionMetadata: self._num_computed_tokens_cache = self.seq_lens - query_lens return self._num_computed_tokens_cache + def token_to_req_indices(self, buffer: torch.Tensor) -> torch.Tensor: + """Build or reuse the per-token request index mapping.""" + num_tokens = self.num_actual_tokens + if self._token_to_req_indices_cache is not None: + assert self._token_to_req_indices_cache.device == buffer.device + assert self._token_to_req_indices_cache.dtype == torch.int32 + assert self._token_to_req_indices_cache.shape[0] >= num_tokens + return self._token_to_req_indices_cache[:num_tokens] + + starts = np.asarray(self.query_start_loc_cpu, dtype=np.int32) + query_lens = np.diff(starts) + token_to_req_indices = np.repeat( + np.arange(query_lens.shape[0], dtype=np.int32), query_lens + ) + num_mapped_tokens = token_to_req_indices.shape[0] + assert buffer.shape[0] >= max(num_mapped_tokens, num_tokens) + # copy from CPU to GPU + buffer[:num_mapped_tokens].copy_( + np_to_pinned_tensor(token_to_req_indices), non_blocking=True + ) + if num_mapped_tokens < num_tokens: + buffer[num_mapped_tokens:num_tokens].zero_() + self._token_to_req_indices_cache = buffer[: max(num_mapped_tokens, num_tokens)] + return self._token_to_req_indices_cache[:num_tokens] + # TODO(lucas): remove once we have FULL-CG spec-decode support def unpadded( self, num_actual_tokens: int, num_actual_reqs: int diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index 6b2d202d3f1..a5735bf313f 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/vllm/v1/attention/backends/cpu_attn.py @@ -437,27 +437,24 @@ def _riscv_supports_rvv() -> bool: The RVV path is compiled whenever __riscv_v_min_vlen is defined, so we check that at least one supported zvlb is advertised. """ + # The C++ compile-time check is the ground truth: it knows which + # VLEN the binary was actually compiled for. The cpuinfo check + # below is only a fast-path shortcut. + try: + import torch + + if torch.ops._C.cpu_attn_has_isa("rvv"): + return True + except Exception: + pass + + # Fallback: check /proc/cpuinfo for zvl128b/zvl256b. try: with open("/proc/cpuinfo") as f: cpuinfo = f.read() except OSError: return False - # If VLEN >= 512 is detected, the RVV kernel was not compiled. - if any(f"zvl{n}b" in cpuinfo for n in (512, 1024)): - return False - - # zvl128b or zvl256b explicitly advertised -> RVV kernel available. - if any(f"zvl{n}b" in cpuinfo for n in (128, 256)): - return True - - # No zvlb flag at all (e.g. some hardware reports zve* without - # a VLEN hint). Delegate to the C++ compile-time check instead. - try: - import torch - - return torch.ops._C.cpu_attn_has_isa("rvv") - except Exception: - return False + return any(f"zvl{n}b" in cpuinfo for n in (128, 256)) def _get_attn_isa( diff --git a/vllm/v1/attention/backends/hpc_attn.py b/vllm/v1/attention/backends/hpc_attn.py index 4a3a4383e2e..8c6dcfe5168 100644 --- a/vllm/v1/attention/backends/hpc_attn.py +++ b/vllm/v1/attention/backends/hpc_attn.py @@ -31,8 +31,6 @@ from vllm.v1.attention.backend import ( ) from vllm.v1.attention.backends.utils import ( KVCacheLayoutType, - get_per_layer_parameters, - infer_global_hyperparameters, split_decodes_and_prefills, ) from vllm.v1.kv_cache_interface import AttentionSpec @@ -86,9 +84,23 @@ class HpcAttnMetadata(AttentionMetadata): hpc_prefill_q_scale: torch.Tensor | None = None """FP8 per-token-per-head Q scale for prefill (from RopeNorm).""" hpc_decode_q_scale: torch.Tensor | None = None - """FP8 per-token-per-head Q scale for decode (from RopeNorm).""" + """FP8 per-token-per-head Q scale for decode (persistent buffer). + shape = [max_decode_tokens, num_q_heads], contiguous. + Only the first num_decode_q_scale_tokens rows are valid.""" hpc_split_k_flag: torch.Tensor | None = None - """Split-K flag tensor for FP8 decode (from RopeNorm).""" + """Split-K flag tensor for FP8 decode (persistent buffer). + shape = [max_num_seqs, num_kv_heads], int32.""" + + # --- MTP (Multi-Token Prediction) fields --- + decode_query_len: int = 1 + """Number of query tokens per decode request. + 1 for standard decoding, mtp+1 for speculative decoding (2 or 3).""" + qo_indptr_decode: torch.Tensor | None = None + """Cumulative query offsets for decode requests (GPU tensor). + shape = [num_decodes + 1]. Only set when decode_query_len > 1. + e.g. 3 requests with dql=2: [0, 2, 4, 6].""" + task_map: torch.Tensor | None = None + """Used for HPC dynamic schedule attention""" class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]): @@ -105,20 +117,40 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]): device: torch.device, ): super().__init__(kv_cache_spec, layer_names, vllm_config, device) - self.model_config = vllm_config.model_config - self.cache_config = vllm_config.cache_config + import hpc - self.num_qo_heads = self.model_config.get_num_attention_heads( - vllm_config.parallel_config - ) self.num_kv_heads = kv_cache_spec.num_kv_heads - self.head_dim = kv_cache_spec.head_size - self.page_size = kv_cache_spec.block_size + self.hpc_dynamic_sched_attn_min_split_len = 1024 - self.cache_dtype = self.cache_config.cache_dtype + # MTP constraint: HPC decode kernel only supports mtp in {0, 1, 2, 3} + spec_config = vllm_config.speculative_config + if ( + spec_config is not None + and spec_config.num_speculative_tokens is not None + and spec_config.num_speculative_tokens > 3 + ): + raise ValueError( + f"HPC attention only supports up to 3 speculative tokens " + f"(mtp ∈ {{0, 1, 2, 3}}), got " + f"num_speculative_tokens={spec_config.num_speculative_tokens}. " + f"Please reduce num_speculative_tokens or use a different " + f"attention backend." + ) - self.global_hyperparameters = infer_global_hyperparameters( - get_per_layer_parameters(vllm_config, layer_names, HpcAttentionImpl) + # Dynamic decode threshold for MTP support. + # _init_reorder_batch_threshold computes: + # no spec_config → threshold=1 (unchanged) + # with spec_config → threshold=1+num_speculative_tokens + self._init_reorder_batch_threshold( + reorder_batch_threshold=1, + supports_spec_as_decode=True, + ) + + self.task_map = hpc.get_attention_decode_task_workspace( + vllm_config.scheduler_config.max_num_seqs, + vllm_config.model_config.max_model_len or 4096, + self.num_kv_heads, + min_process_len=self.hpc_dynamic_sched_attn_min_split_len, ) @override # type: ignore[misc] @@ -128,6 +160,13 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]): vllm_config: VllmConfig, kv_cache_spec: AttentionSpec, ) -> AttentionCGSupport: + spec_config = vllm_config.speculative_config + if ( + spec_config is not None + and spec_config.num_speculative_tokens is not None + and spec_config.num_speculative_tokens > 0 + ): + return AttentionCGSupport.UNIFORM_BATCH return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE def build( @@ -143,7 +182,8 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]): split_decodes_and_prefills( common_attn_metadata, decode_threshold=self.reorder_batch_threshold, - require_uniform=False, + # MTP requires uniform query lengths across decode requests + require_uniform=(self.reorder_batch_threshold > 1), ) ) @@ -152,7 +192,16 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]): slot_mapping = common_attn_metadata.slot_mapping max_query_len = common_attn_metadata.max_query_len + # Compute decode_query_len (tokens per decode request). + # Non-MTP: 1, MTP: mtp+1 (2 or 3). + if num_decodes > 0 and num_decode_tokens > num_decodes: + decode_query_len = num_decode_tokens // num_decodes + else: + decode_query_len = 1 + + seq_lens_decode = None qo_indptr = None + qo_indptr_decode = None if num_prefills > 0: qo_indptr_cpu = common_attn_metadata.query_start_loc_cpu prefill_start = num_decodes @@ -161,6 +210,21 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]): ) qo_indptr = qo_indptr_prefill_cpu.to(self.device, non_blocking=True) + if num_decodes > 0: + seq_lens_decode = seq_lens[:num_decodes] + # block_table is per-request, indexed by num_decodes (not tokens) + qo_indptr_decode = common_attn_metadata.query_start_loc[: num_decodes + 1] + import hpc + + hpc.assign_attention_decode_task( + seq_lens_decode, + self.task_map, + self.num_kv_heads, + decode_query_len, + new_kv_included=True, + min_process_len=self.hpc_dynamic_sched_attn_min_split_len, + ) + return HpcAttnMetadata( num_actual_tokens=num_actual_tokens, num_decodes=num_decodes, @@ -176,6 +240,9 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]): hpc_prefill_q_scale=None, hpc_decode_q_scale=None, hpc_split_k_flag=None, + decode_query_len=decode_query_len, + qo_indptr_decode=qo_indptr_decode, + task_map=self.task_map, ) @@ -192,6 +259,7 @@ class HpcAttentionBackend(AttentionBackend): ] supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ "auto", + "bfloat16", "fp8_e4m3", ] @@ -323,6 +391,13 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]): self.supports_quant_query_input = False self.splitk = True + import hpc + + if self.use_fp8: + self._quant_type = hpc.QuantType.QPERTOKEN_PERHEAD_KPERTENSOR_VPERTENSOR + else: + self._quant_type = None + def forward( self, layer: torch.nn.Module, @@ -417,6 +492,7 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]): block_table_prefill, seq_lens_prefill, max_seqlens, + quant_type=self._quant_type, output=output_prefill, ) else: @@ -439,6 +515,8 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]): q_decode = query[:num_decode_tokens] output_decode = output[:num_decode_tokens] + mtp = attn_metadata.decode_query_len - 1 + if self.use_fp8: hpc.attention_decode_fp8( q_decode, @@ -449,8 +527,14 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]): hpc_decode_q_scale, k_scale, v_scale, + mtp=mtp, + # MTP: split_flag from rope_norm is unavailable + # when using prefill-mode kernel; let HPC decide. + # splitk=(self.splitk if mtp == 0 else True), new_kv_included=True, + quant_type=self._quant_type, splitk=self.splitk, + task_map=attn_metadata.task_map, split_flag=hpc_split_k_flag, output=output_decode, ) @@ -461,6 +545,7 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]): kv_cache[:, 1], block_table_decode, num_seq_kvcache, + mtp=mtp, output=output_decode, new_kv_included=True, splitk=self.splitk, diff --git a/vllm/v1/attention/backends/linear_attn.py b/vllm/v1/attention/backends/linear_attn.py index b2ca151986c..9cdcf0e30e7 100644 --- a/vllm/v1/attention/backends/linear_attn.py +++ b/vllm/v1/attention/backends/linear_attn.py @@ -4,7 +4,7 @@ from dataclasses import dataclass import torch -from vllm.config import VllmConfig +from vllm.config import CompilationConfig, VllmConfig from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -12,6 +12,7 @@ from vllm.v1.attention.backend import ( CommonAttentionMetadata, ) from vllm.v1.attention.backends.utils import ( + PAD_SLOT_ID, mamba_get_block_table_tensor, split_decodes_and_prefills, ) @@ -91,3 +92,213 @@ class LinearAttentionMetadataBuilder(AttentionMetadataBuilder[LinearAttentionMet state_indices_tensor=state_indices_tensor, ) return attn_metadata + + +class BailingLinearAttentionBackend(LinearAttentionBackend): + @staticmethod + def get_name() -> str: + return "BAILING_LINEAR_ATTN" + + @staticmethod + def get_builder_cls() -> type["BailingLinearAttentionMetadataBuilder"]: + return BailingLinearAttentionMetadataBuilder + + +@dataclass +class BailingLinearAttentionMetadata(LinearAttentionMetadata): + state_indices_tensor_d: torch.Tensor | None = None + state_indices_tensor_p: torch.Tensor | None = None + num_accepted_tokens: torch.Tensor | None = None + query_start_loc_d: torch.Tensor | None = None + + +class BailingLinearAttentionMetadataBuilder(LinearAttentionMetadataBuilder): + supports_spec_decode_metadata = True + supports_update_block_table: bool = False + + @classmethod + def get_cudagraph_support( + cls, + vllm_config: VllmConfig, + kv_cache_spec: AttentionSpec, + ) -> AttentionCGSupport: + return AttentionCGSupport.UNIFORM_BATCH + + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ): + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + self.compilation_config: CompilationConfig = vllm_config.compilation_config + self.num_spec_tokens: int = vllm_config.num_speculative_tokens + self.use_spec_decode: bool = self.num_spec_tokens > 0 + self.decode_cudagraph_max_bs: int = vllm_config.scheduler_config.max_num_seqs + if self.compilation_config.max_cudagraph_capture_size is not None: + self.decode_cudagraph_max_bs = min( + self.decode_cudagraph_max_bs, + self.compilation_config.max_cudagraph_capture_size, + ) + self.decode_state_indices_tensor: torch.Tensor = torch.empty( + (self.decode_cudagraph_max_bs, 1 + self.num_spec_tokens), + dtype=torch.int32, + device=device, + ) + self.decode_legacy_state_indices_tensor: torch.Tensor = torch.empty( + (self.decode_cudagraph_max_bs,), + dtype=torch.int32, + device=device, + ) + self.decode_query_start_loc: torch.Tensor = torch.empty( + (self.decode_cudagraph_max_bs + 1,), + dtype=torch.int32, + device=device, + ) + self.decode_num_accepted_tokens: torch.Tensor = torch.empty( + (self.decode_cudagraph_max_bs,), + dtype=torch.int32, + device=device, + ) + self._init_reorder_batch_threshold(1, self.use_spec_decode) + + def build_for_cudagraph_capture( + self, + common_attn_metadata: CommonAttentionMetadata, + ) -> BailingLinearAttentionMetadata: + num_accepted_tokens = None + if self.use_spec_decode: + assert common_attn_metadata.max_query_len <= 1 + self.num_spec_tokens, ( + "Bailing linear attention only supports speculative decoding " + "with query length <= 1 + number of speculative tokens." + ) + num_accepted_tokens = torch.diff(common_attn_metadata.query_start_loc) + return self.build( + common_prefix_len=0, + common_attn_metadata=common_attn_metadata, + num_accepted_tokens=num_accepted_tokens, + ) + + def build( # type: ignore[override] + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + *, + num_accepted_tokens: torch.Tensor | None = None, + num_decode_draft_tokens_cpu: torch.Tensor | None = None, + ) -> BailingLinearAttentionMetadata: + query_start_loc = common_attn_metadata.query_start_loc + seq_lens = common_attn_metadata.seq_lens + num_reqs = common_attn_metadata.num_reqs + use_spec_decode = self.use_spec_decode and num_accepted_tokens is not None + + state_indices_tensor = mamba_get_block_table_tensor( + common_attn_metadata.block_table_tensor, + common_attn_metadata.seq_lens, + self.kv_cache_spec, + self.vllm_config.cache_config.mamba_cache_mode, + ) + if state_indices_tensor.dim() == 1: + state_indices_tensor = state_indices_tensor.unsqueeze(-1) + + decode_threshold = self.reorder_batch_threshold if use_spec_decode else 1 + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( + split_decodes_and_prefills( + common_attn_metadata, + decode_threshold=decode_threshold, + ) + ) + state_indices_tensor_d, state_indices_tensor_p = torch.split( + state_indices_tensor, + [num_decodes, num_prefills], + dim=0, + ) + state_indices_tensor_p = state_indices_tensor_p[:, 0] + + query_start_loc_d = None + if use_spec_decode: + assert num_accepted_tokens is not None + state_indices_tensor_d = state_indices_tensor_d[ + :, : 1 + self.num_spec_tokens + ] + query_start_loc_d = query_start_loc[: num_decodes + 1] + num_accepted_tokens = num_accepted_tokens[:num_decodes] + else: + state_indices_tensor_d = state_indices_tensor_d[:, 0] + num_accepted_tokens = None + + legacy_state_indices_tensor = state_indices_tensor[:, 0] + cudagraph_mode = self.compilation_config.cudagraph_mode + use_full_cudagraph = ( + cudagraph_mode is not None and cudagraph_mode.has_full_cudagraphs() + ) + if ( + num_prefills == 0 + and num_decodes <= self.decode_cudagraph_max_bs + and use_full_cudagraph + ): + padded_bs = num_reqs + is_padded_decode = seq_lens[:num_decodes] == 0 + if state_indices_tensor_d.dim() > 1: + state_indices_tensor_d = torch.where( + is_padded_decode.unsqueeze(1), + torch.full_like(state_indices_tensor_d, PAD_SLOT_ID), + state_indices_tensor_d, + ) + self.decode_state_indices_tensor[:num_decodes].copy_( + state_indices_tensor_d, + non_blocking=True, + ) + state_indices_tensor_d = self.decode_state_indices_tensor[:padded_bs] + state_indices_tensor_d[num_decodes:] = PAD_SLOT_ID + + self.decode_legacy_state_indices_tensor[:num_decodes].copy_( + torch.where( + is_padded_decode, + torch.full_like( + legacy_state_indices_tensor[:num_decodes], + PAD_SLOT_ID, + ), + legacy_state_indices_tensor[:num_decodes], + ), + non_blocking=True, + ) + legacy_state_indices_tensor = self.decode_legacy_state_indices_tensor[ + :padded_bs + ] + legacy_state_indices_tensor[num_decodes:] = PAD_SLOT_ID + if state_indices_tensor_d.dim() == 1: + state_indices_tensor_d = legacy_state_indices_tensor + + if use_spec_decode and num_accepted_tokens is not None: + assert query_start_loc_d is not None + self.decode_query_start_loc[: num_decodes + 1].copy_( + query_start_loc_d, + non_blocking=True, + ) + decode_num_query_tokens = query_start_loc_d[-1] + query_start_loc_d = self.decode_query_start_loc[: padded_bs + 1] + query_start_loc_d[num_decodes + 1 :] = decode_num_query_tokens + + self.decode_num_accepted_tokens[:num_decodes].copy_( + num_accepted_tokens, + non_blocking=True, + ) + num_accepted_tokens = self.decode_num_accepted_tokens[:padded_bs] + num_accepted_tokens[num_decodes:] = 1 + + return BailingLinearAttentionMetadata( + num_prefills=num_prefills, + num_prefill_tokens=num_prefill_tokens, + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, + query_start_loc=query_start_loc, + seq_lens=seq_lens, + state_indices_tensor=legacy_state_indices_tensor, + state_indices_tensor_d=state_indices_tensor_d, + state_indices_tensor_p=state_indices_tensor_p, + num_accepted_tokens=num_accepted_tokens, + query_start_loc_d=query_start_loc_d, + ) diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index e6a64ee85f8..977a42c6fef 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -196,6 +196,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): kv_cache_dtype_str = getattr(vllm_config.cache_config, "cache_dtype", "auto") if kv_cache_dtype_str in ("fp8", "fp8_e4m3", "fp8_e5m2"): kv_cache_dtype_str = "fp8" + q_dtype = dtypes.fp8 else: kv_cache_dtype_str = "bf16" kv_dtype = dtypes.d_dtypes.get(kv_cache_dtype_str, dtypes.bf16) diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index ac722dca9fc..ca5be3a3020 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -23,6 +23,7 @@ from vllm.v1.kv_cache_interface import ( KVCacheSpec, MLAAttentionSpec, SlidingWindowMLASpec, + get_kv_quant_mode, ) # DeepseekV4 decode layer types, keyed by compress_ratio. Each type has a distinct @@ -91,6 +92,7 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase): cache_dtype_str=self.cache_config.cache_dtype, alignment=576 if uses_fp8_ds_mla_layout else None, model_version="deepseek_v4", + kv_quant_mode=get_kv_quant_mode(self.cache_config.cache_dtype), ) def forward(self): ... @@ -399,7 +401,6 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): For prefill, we use chunked prefill to align with the indexer's chunking. """ - num_reqs = common_attn_metadata.num_reqs seq_lens = common_attn_metadata.seq_lens seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound query_start_loc = common_attn_metadata.query_start_loc @@ -416,10 +417,9 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): # NOTE: Ensure all metadata tensors maintain fixed memory addresses # for CUDA graph compatibility. - query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] - x = torch.repeat_interleave(torch.arange(num_reqs), query_lens).pin_memory() - token_to_req_indices = self.token_to_req_indices[: x.shape[0]] - token_to_req_indices.copy_(x, non_blocking=True) + token_to_req_indices = common_attn_metadata.token_to_req_indices( + self.token_to_req_indices + ) is_valid_token = self.is_valid_token[: slot_mapping.shape[0]] is_valid_token.copy_(slot_mapping >= 0) diff --git a/vllm/v1/attention/ops/triton_decode_attention.py b/vllm/v1/attention/ops/triton_decode_attention.py index dbe3c5705de..6aec1db2e59 100644 --- a/vllm/v1/attention/ops/triton_decode_attention.py +++ b/vllm/v1/attention/ops/triton_decode_attention.py @@ -133,7 +133,7 @@ def _fwd_kernel_stage1( + offs_n // PAGE_SIZE, mask=offs_n < split_kv_end, other=0, - ) + ).to(tl.int64) # page_number * page stride overflows int32 kv_in_page = offs_n % PAGE_SIZE offs_buf_k = ( (kv_page_number * stride_buf_kpbs + kv_in_page * stride_buf_kbs)[ @@ -375,7 +375,7 @@ def _fwd_grouped_kernel_stage1( mask=offs_n < split_kv_end, other=0, cache_modifier=".ca", - ) + ).to(tl.int64) # page_number * page stride overflows int32 kv_off_k = ( kv_page_number * stride_buf_kpbs + (offs_n % PAGE_SIZE) * stride_buf_kbs ) diff --git a/vllm/v1/attention/ops/triton_unified_attention.py b/vllm/v1/attention/ops/triton_unified_attention.py index a13ae96a7a9..93622957b55 100644 --- a/vllm/v1/attention/ops/triton_unified_attention.py +++ b/vllm/v1/attention/ops/triton_unified_attention.py @@ -254,12 +254,19 @@ def kernel_unified_attention( # Per-(token, head) scale caches: used iff KV_QUANT_MODE in {2, 3}. k_scale_cache_ptr=None, v_scale_cache_ptr=None, - stride_ks_blk: tl.int64 = None, - stride_ks_slot: tl.int64 = None, - stride_ks_head: tl.int64 = None, - stride_vs_blk: tl.int64 = None, - stride_vs_slot: tl.int64 = None, - stride_vs_head: tl.int64 = None, + # ``tl.int64`` cannot be combined with a ``None`` default — Triton's JIT + # rejects ``Optional[tl.int64]`` / ``tl.int64 | None`` at trace time, and + # plain ``tl.int64 = None`` raises ``TypeError: 'NoneType' object cannot + # be interpreted as an integer`` when callers omit these arguments. + # ``int | None`` is the only annotation that lets the wrapper pass + # ``None`` here so Triton can skip materialising the strides when the + # ``USE_PER_TOKEN_HEAD_SCALES`` branch is dead. + stride_ks_blk: int | None = None, + stride_ks_slot: int | None = None, + stride_ks_head: int | None = None, + stride_vs_blk: int | None = None, + stride_vs_slot: int | None = None, + stride_vs_head: int | None = None, # KV cache quantization mode handled inside this kernel via constexpr # branches: NONE (0), FP8_PER_TENSOR (1), INT8_PER_TOKEN_HEAD (2), # FP8_PER_TOKEN_HEAD (3). Sub-byte INT4 (4) uses its own @@ -283,7 +290,10 @@ def kernel_unified_attention( # original (causal AND SW) OR mm_prefix behavior for all other models. MM_PREFIX_CLAMP_SW: tl.constexpr = False, ): - USE_PER_TOKEN_HEAD_SCALES: tl.constexpr = KV_QUANT_MODE >= 2 + # Per-(token, head) scale caches: used iff KV_QUANT_MODE in {2, 3}. + USE_PER_TOKEN_HEAD_SCALES: tl.constexpr = (KV_QUANT_MODE >= 2) and ( + KV_QUANT_MODE <= 3 + ) USE_FP8_Q_DESCALE: tl.constexpr = KV_QUANT_MODE == 1 and Q_IS_FP8 if USE_TD: @@ -1041,9 +1051,9 @@ def unified_attention( # The kernel signature is the same for 2D and 3D — only the launch # grid + a handful of constexpr toggles differ. Per-token-head scale - # caches and their strides are required arguments; non-per-token-head - # modes pass dummy zeros (the code path is dead-code eliminated by - # the ``USE_PER_TOKEN_HEAD_SCALES`` constexpr branch in the kernel). + # caches and their strides are passed as ``None`` when the + # ``USE_PER_TOKEN_HEAD_SCALES`` branch is dead so Triton can skip + # materialising those arguments and the associated registers. if use_per_token_head_scales: ks_strides = k_scale_cache.stride() vs_strides = v_scale_cache.stride() @@ -1052,16 +1062,15 @@ def unified_attention( k_scale_ptr = k_scale_cache v_scale_ptr = v_scale_cache else: - ks_blk = ks_slot = ks_head = 0 - vs_blk = vs_slot = vs_head = 0 - # Pass the K cache as a stand-in pointer; never dereferenced. - k_scale_ptr = k - v_scale_ptr = v - # 3D needs real segm tensors; 2D never touches them but Triton wants - # a non-null pointer. Reuse ``out`` as the placeholder. - segm_output_ptr = softmax_segm_output if use_3d else out - segm_max_ptr = softmax_segm_max if use_3d else out - segm_expsum_ptr = softmax_segm_expsum if use_3d else out + ks_blk = ks_slot = ks_head = None + vs_blk = vs_slot = vs_head = None + k_scale_ptr = None + v_scale_ptr = None + # 3D needs real segm tensors; 2D never touches them. Pass ``None`` in + # 2D mode so Triton can skip materialising these pointer arguments. + segm_output_ptr = softmax_segm_output if use_3d else None + segm_max_ptr = softmax_segm_max if use_3d else None + segm_expsum_ptr = softmax_segm_expsum if use_3d else None num_segments = num_par_softmax_segments if use_3d else 1 grid: tuple[Any, ...] diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 9ce1d94ef3c..f75bfae2b54 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1585,8 +1585,12 @@ class Scheduler(SchedulerInterface): scheduled_spec_token_ids = ( scheduler_output.scheduled_spec_decode_tokens.get(req_id) ) - if scheduled_spec_token_ids and ( - generated_token_ids or self.num_sampled_tokens_per_step == 0 + # Skip a stale frame still pending discard (async_tokens_to_discard + # > 0): its pre-reset rejection count would underflow the counters. + if ( + scheduled_spec_token_ids + and (generated_token_ids or self.num_sampled_tokens_per_step == 0) + and request.async_tokens_to_discard == 0 ): num_draft_tokens = len(scheduled_spec_token_ids) num_sampled = self.num_sampled_tokens_per_step diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index d5cf1050ca4..bcb441e7564 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -1434,6 +1434,12 @@ class DPLBAsyncMPClient(DPAsyncMPClient): # Increment local waiting count for better balancing between stats # updates from the coordinator (which happen every 100ms). current_counts[eng_index][0] += self.client_count + # Rotate the scan start so that ties (equal scores, e.g. right + # after a coordinator stats reset when engines look equally loaded) + # don't systematically favor the same engine. This removes the + # fixed tie-break bias without affecting load-aware decisions when + # scores actually differ. + self.eng_start_index = (self.eng_start_index + 1) % num_engines chosen_engine = self.core_engines[eng_index] # Record which engine is chosen for this request, to handle aborts. diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 4d27e308e88..756c5f3b371 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -1468,6 +1468,26 @@ class SpecDecodeBaseProposer: "Sharing target model embedding weights with the draft model." ) + if share_embeddings: + draft_embed = self.model.model.embed_tokens + # Only share when both models use the same embedding width. + # Guard with isinstance so non-Tensor weights (e.g. in tests) + # are not affected — mirrors the weight-equality check above. + if isinstance(target_embed_tokens.weight, torch.Tensor) and isinstance( + draft_embed.weight, torch.Tensor + ): + target_dim = target_embed_tokens.weight.shape[-1] + draft_dim = draft_embed.weight.shape[-1] + if target_dim != draft_dim: + share_embeddings = False + logger.info( + "Target embedding dim (%d) differs from draft " + "embedding dim (%d). Keeping separate embedding " + "weights.", + target_dim, + draft_dim, + ) + if share_embeddings: if hasattr(self.model.model, "embed_tokens"): del self.model.model.embed_tokens diff --git a/vllm/v1/spec_decode/ngram_proposer_gpu.py b/vllm/v1/spec_decode/ngram_proposer_gpu.py index a08b341e80a..ed544bb27c1 100644 --- a/vllm/v1/spec_decode/ngram_proposer_gpu.py +++ b/vllm/v1/spec_decode/ngram_proposer_gpu.py @@ -468,7 +468,7 @@ class NgramProposerGPU: def update_scheduler_for_invalid_drafts( - num_valid_draft_tokens_event: torch.Event, + num_valid_draft_tokens_event: torch.cuda.Event, num_valid_draft_tokens_cpu: torch.Tensor, scheduler_output: "SchedulerOutput", req_id_to_index: dict[str, int], @@ -643,7 +643,7 @@ def _sync_num_tokens( def copy_num_valid_draft_tokens( num_valid_draft_tokens_cpu: torch.Tensor, num_valid_draft_tokens_copy_stream: torch.cuda.Stream, - num_valid_draft_tokens_event: torch.Event, + num_valid_draft_tokens_event: torch.cuda.Event, num_valid_draft_tokens: torch.Tensor | None, batch_size: int, ) -> None: diff --git a/vllm/v1/worker/cpu/shm.py b/vllm/v1/worker/cpu/shm.py index deec52e44ba..bd1f96c71ed 100644 --- a/vllm/v1/worker/cpu/shm.py +++ b/vllm/v1/worker/cpu/shm.py @@ -48,6 +48,7 @@ def get_memory_info(*args: Any, **kwargs: Any) -> tuple[int, int]: torch.Event = _EventPlaceholder +torch.cuda.Event = _EventPlaceholder torch.cuda.Stream = _StreamPlaceholder torch.cuda.set_stream = noop torch.cuda.current_stream = lambda *args, **kwargs: _StreamPlaceholder() diff --git a/vllm/v1/worker/gpu/async_utils.py b/vllm/v1/worker/gpu/async_utils.py index a9ad16b1520..b3d6f5e4d90 100644 --- a/vllm/v1/worker/gpu/async_utils.py +++ b/vllm/v1/worker/gpu/async_utils.py @@ -24,7 +24,7 @@ class AsyncOutput(AsyncModelRunnerOutput): self.model_runner_output = model_runner_output self.sampler_output = sampler_output self.num_sampled_tokens = num_sampled_tokens - self.copy_event = torch.Event() + self.copy_event = torch.cuda.Event() with stream(copy_stream, main_stream): copy_stream.wait_stream(main_stream) @@ -81,7 +81,7 @@ class AsyncPoolingOutput(AsyncModelRunnerOutput): self.model_runner_output = model_runner_output self.pooler_output = pooler_output self.is_valid = is_valid - self.copy_event = torch.Event() + self.copy_event = torch.cuda.Event() with stream(copy_stream, main_stream): copy_stream.wait_stream(main_stream) diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index aecce3c575c..dfa2daacf9f 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -12,6 +12,7 @@ from vllm.config import ( get_layers_from_vllm_config, set_current_vllm_config, ) +from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.multimodal.inputs import MultiModalFeatureSpec @@ -37,6 +38,8 @@ from vllm.v1.worker.utils import ( prepare_kernel_block_sizes, ) +logger = init_logger(__name__) + @dataclass(frozen=True) class AttentionCGSupportInfo: @@ -368,6 +371,14 @@ def _reshape_kv_cache( kernel_block_sizes=kernel_block_sizes, cache_dtype=cache_dtype, ) + elif has_attn and kv_cache_config is not None: + _align_mixed_attention_kv_cache_views( + attn_groups=attn_groups, + kv_caches=kv_caches, + kernel_block_sizes=kernel_block_sizes, + cache_dtype=cache_dtype, + kv_cache_config=kv_cache_config, + ) # Map any sharing layers to their target layer's KV cache. for layer_name, target_layer_name in shared_kv_cache_layers.items(): @@ -376,6 +387,77 @@ def _reshape_kv_cache( return kv_caches +def _align_mixed_attention_kv_cache_views( + attn_groups: Iterable[AttentionGroup], + kv_caches: dict[str, Any], + kernel_block_sizes: list[int], + cache_dtype: str, + kv_cache_config: KVCacheConfig, +) -> None: + """Align shared attention KV views when backends disagree on layout. + + Encoder-decoder models can share one raw allocation between decoder + self-attention (K/V-first ROCM_ATTN, block dim 1) and cross-attention + (blocks-first backends, block dim 0). Keep the physical storage in the + K/V-first layout expected by ROCM_ATTN, and restride the blocks-first + logical views so block IDs address the same bytes. + """ + block_dims_by_layer: dict[str, int] = {} + for group in attn_groups: + kv_cache_spec = group.kv_cache_spec + if not isinstance(kv_cache_spec, AttentionSpec): + continue + if group.kv_cache_group_id >= len(kernel_block_sizes): + continue + block_dim = group.backend.get_kv_cache_block_dim( + kernel_block_sizes[group.kv_cache_group_id], + kv_cache_spec.num_kv_heads, + kv_cache_spec.head_size, + cache_dtype_str=cache_dtype, + ) + for layer_name in group.layer_names: + if layer_name in kv_caches: + block_dims_by_layer[layer_name] = block_dim + + for kv_tensor in kv_cache_config.kv_cache_tensors: + if kv_tensor.block_stride > 0: + continue + shared_block_dims = { + block_dims_by_layer[layer_name] + for layer_name in kv_tensor.shared_by + if layer_name in block_dims_by_layer + } + if 0 not in shared_block_dims or 1 not in shared_block_dims: + continue + + for layer_name in kv_tensor.shared_by: + if block_dims_by_layer.get(layer_name) == 0: + _restride_blocks_first_kv_cache_to_kv_first_storage( + kv_caches[layer_name] + ) + + +def _restride_blocks_first_kv_cache_to_kv_first_storage( + kv_cache: torch.Tensor, +) -> None: + assert kv_cache.ndim >= 3 + assert kv_cache.shape[1] == 2 + page_size = kv_cache.shape[2:].numel() + num_blocks = kv_cache.shape[0] + expected_tail_stride = torch.empty(kv_cache.shape[2:]).stride() + if kv_cache.stride()[2:] != expected_tail_stride: + logger.warning_once( + "Skipping mixed KV-cache layout alignment for a non-NHD " + "blocks-first attention view with stride %s.", + kv_cache.stride(), + ) + return + kv_cache.as_strided_( + size=kv_cache.shape, + stride=(page_size, num_blocks * page_size, *expected_tail_stride), + ) + + def _update_hybrid_attention_layout( attn_groups: Iterable[AttentionGroup], kv_caches: dict[str, Any], diff --git a/vllm/v1/worker/gpu/pp_utils.py b/vllm/v1/worker/gpu/pp_utils.py index 00ff95b6dac..9b1786cbe4c 100644 --- a/vllm/v1/worker/gpu/pp_utils.py +++ b/vllm/v1/worker/gpu/pp_utils.py @@ -9,6 +9,7 @@ import numpy as np import torch from vllm.distributed.parallel_state import get_pp_group +from vllm.platforms import current_platform from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu from vllm.v1.worker.gpu.input_batch import InputBatch @@ -17,7 +18,7 @@ from vllm.v1.worker.gpu.input_batch import InputBatch class PendingRecv: """Per-step slot data for a deferred postprocess on the main stream.""" - event: torch.Event + event: torch.cuda.Event sampled_tokens: torch.Tensor # [num_reqs, max_sample_len] num_sampled: torch.Tensor # [num_reqs] @@ -179,6 +180,10 @@ class PPHandler: return assert sampled_token_ids.dtype == torch.int64 + + if current_platform.is_xpu(): + self.main_stream.synchronize() + with torch.cuda.stream(self.broadcast_stream): self.broadcast_stream.wait_stream(self.main_stream) torch.distributed.broadcast( diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py index 37ca1665937..89d46a75822 100644 --- a/vllm/v1/worker/gpu/spec_decode/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/utils.py @@ -12,7 +12,7 @@ class DraftTokensHandler: def __init__(self, device: torch.device | None = None): self.device = device self.copy_stream = torch.cuda.Stream(device) - self.copy_event = torch.Event() + self.copy_event = torch.cuda.Event() self.req_ids: list[str] = [] self.draft_tokens_np: np.ndarray | None = None diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 3930a07b248..24889f4021c 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -132,6 +132,9 @@ from vllm.v1.attention.backend import ( CommonAttentionMetadata, ) from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder +from vllm.v1.attention.backends.linear_attn import ( + BailingLinearAttentionMetadataBuilder, +) from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder from vllm.v1.attention.backends.utils import ( NULL_BLOCK_ID, @@ -857,7 +860,7 @@ class GPUModelRunner( # N-gram GPU path: async D2H buffer/event for per-request valid draft counts. self._num_valid_draft_tokens: torch.Tensor | None = None self._num_valid_draft_tokens_cpu: torch.Tensor | None = None - self._num_valid_draft_tokens_event: torch.Event | None = None + self._num_valid_draft_tokens_event: torch.cuda.Event | None = None self._num_valid_draft_tokens_copy_stream: torch.cuda.Stream | None = None if ( self.speculative_config is not None @@ -866,7 +869,7 @@ class GPUModelRunner( self._num_valid_draft_tokens_cpu = torch.empty( self.max_num_reqs, dtype=torch.int32, pin_memory=PIN_MEMORY ) - self._num_valid_draft_tokens_event = torch.Event() + self._num_valid_draft_tokens_event = torch.cuda.Event() self._num_valid_draft_tokens_copy_stream = torch.cuda.Stream() self._draft_token_req_ids: list[str] | None = None @@ -2442,9 +2445,16 @@ class GPUModelRunner( extra_attn_metadata_args = {} if use_spec_decode and isinstance( - builder, (Mamba2AttentionMetadataBuilder, GDNAttentionMetadataBuilder) + builder, + ( + Mamba2AttentionMetadataBuilder, + GDNAttentionMetadataBuilder, + BailingLinearAttentionMetadataBuilder, + ), ): - assert ubid is None, "UBatching not supported with GDN yet" + assert ubid is None, ( + "UBatching not supported with GDN or linear attn yet" + ) extra_attn_metadata_args = dict( num_accepted_tokens=self.num_accepted_tokens.gpu[:num_reqs_padded], num_decode_draft_tokens_cpu=self.num_decode_draft_tokens.cpu[ @@ -4517,17 +4527,23 @@ class GPUModelRunner( self._copy_draft_token_ids_to_cpu(scheduler_output) spec_config = self.speculative_config - propose_drafts_after_bookkeeping = False + draft_after_bookkeeping = False if spec_config is not None: # Decide whether to run the drafter or zero out draft tokens. input_fits_in_drafter = self._input_fits_in_drafter( spec_decode_common_attn_metadata ) - use_gpu_toks = ( + # Whether the drafter runs a GPU model forward (and thus carries + # TP/EP/DP collectives), independent of padded-batch timing. + drafter_runs_model_forward = ( spec_config.use_eagle() or spec_config.uses_draft_model() or spec_config.uses_extract_hidden_states() - ) and not spec_config.disable_padded_drafter_batch + ) + use_gpu_toks = ( + drafter_runs_model_forward + and not spec_config.disable_padded_drafter_batch + ) if use_gpu_toks: # EAGLE/DraftModel speculative decoding can use the GPU sampled tokens # as inputs, and does not need to wait for bookkeeping to finish. @@ -4542,19 +4558,23 @@ class GPUModelRunner( sampled_token_ids = sampler_output.sampled_token_ids if input_fits_in_drafter: propose_draft_token_ids(sampled_token_ids) - elif self.valid_sampled_token_count_event is not None: - assert spec_decode_common_attn_metadata is not None - next_token_ids, valid_sampled_tokens_count = ( - self.drafter.prepare_next_token_ids_padded( - sampled_token_ids, - self.requests, - self.input_batch, - self.discard_request_mask.gpu, + else: + if self.valid_sampled_token_count_event is not None: + assert spec_decode_common_attn_metadata is not None + next_token_ids, valid_sampled_tokens_count = ( + self.drafter.prepare_next_token_ids_padded( + sampled_token_ids, + self.requests, + self.input_batch, + self.discard_request_mask.gpu, + ) ) - ) - self._copy_valid_sampled_token_count( - next_token_ids, valid_sampled_tokens_count - ) + self._copy_valid_sampled_token_count( + next_token_ids, valid_sampled_tokens_count + ) + if self.parallel_config.data_parallel_size > 1: + # Prevent hang when DP ranks disagree on input_fits_in_drafter + self.drafter.dummy_run(num_tokens=1) elif ( spec_config.use_ngram_gpu() and not spec_config.disable_padded_drafter_batch @@ -4578,7 +4598,9 @@ class GPUModelRunner( next_token_ids, valid_sampled_tokens_count ) else: - propose_drafts_after_bookkeeping = input_fits_in_drafter + # These drafters consume CPU sampled tokens, so they run + # after bookkeeping. + draft_after_bookkeeping = True if not input_fits_in_drafter: # Zero out draft tokens so the scheduler doesn't schedule @@ -4610,10 +4632,25 @@ class GPUModelRunner( scheduler_output.total_num_scheduled_tokens, ) - if propose_drafts_after_bookkeeping: + if draft_after_bookkeeping: # ngram and other speculative decoding methods use the sampled # tokens on the CPU, so they are run after bookkeeping. - propose_draft_token_ids(valid_sampled_token_ids) + if input_fits_in_drafter: + propose_draft_token_ids(valid_sampled_token_ids) + elif ( + drafter_runs_model_forward + and self.parallel_config.data_parallel_size > 1 + ): + # Prevent hang when DP ranks disagree on input_fits_in_drafter + assert isinstance( + self.drafter, + EagleProposer + | DFlashProposer + | DraftModelProposer + | ExtractHiddenStatesProposer + | Gemma4Proposer, + ) + self.drafter.dummy_run(num_tokens=1) # Finalize KV connector (wait_for_save + clear metadata) after # draft model runs. Deferred from target model forward to allow @@ -7154,7 +7191,12 @@ class GPUModelRunner( layer_cache_dtype_str = ( "auto" if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE - else self.cache_config.cache_dtype + else getattr( + kv_cache_spec, + "cache_dtype_str", + None, + ) + or self.cache_config.cache_dtype ) kv_cache_shape = attn_backend.get_kv_cache_shape( kernel_num_blocks, diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 28efee3dee8..03433ed7524 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -55,7 +55,7 @@ from vllm.multimodal.video import ( PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES, PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES, PYNVVIDEOCODEC_MAX_RETAINED_DECODERS, - PYNVVIDEOCODEC_VIDEO_BACKEND, + VIDEO_LOADER_REGISTRY, ) from vllm.platforms import current_platform from vllm.profiler.wrapper import CudaProfilerWrapper, TorchProfilerWrapper @@ -584,15 +584,15 @@ class Worker(WorkerBase): ) @staticmethod - def _uses_pynvvideocodec_video_backend(mm_config) -> bool: + def _uses_gpu_video_backend(mm_config) -> bool: video_kwargs = mm_config.media_io_kwargs.get("video", {}) video_loader_backend = ( video_kwargs.get("video_backend") or envs.VLLM_VIDEO_LOADER_BACKEND ) codec_backend = video_kwargs.get("backend") - return ( - video_loader_backend == PYNVVIDEOCODEC_VIDEO_BACKEND - or codec_backend == PYNVVIDEOCODEC_VIDEO_BACKEND + return VIDEO_LOADER_REGISTRY.backend_requires_gpu(video_loader_backend) or ( + codec_backend is not None + and VIDEO_LOADER_REGISTRY.backend_requires_gpu(codec_backend) ) def _reserve_mm_ipc_gpu_memory(self, available_kv_cache_memory_bytes: int) -> int: @@ -623,7 +623,7 @@ class Worker(WorkerBase): ) decoder_reserved_bytes = ( num_api_servers * per_server_decoder_bytes - if self._uses_pynvvideocodec_video_backend(mm_config) + if self._uses_gpu_video_backend(mm_config) else 0 ) reserved_bytes = raw_frame_reserved_bytes + decoder_reserved_bytes diff --git a/vllm/v1/worker/xpu_model_runner.py b/vllm/v1/worker/xpu_model_runner.py index 6cdca994da5..82f129b7e60 100644 --- a/vllm/v1/worker/xpu_model_runner.py +++ b/vllm/v1/worker/xpu_model_runner.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from contextlib import contextmanager +from functools import partial import torch @@ -40,14 +41,17 @@ class XPUModelRunnerV2(GPUModelRunnerV2): @contextmanager def _torch_cuda_wrapper(): - # replace cuda APIs with xpu APIs, this should work by default + # Replace cuda APIs with xpu APIs. Each callable gets its own functools.partial + # so it is not the same object as torch.xpu.* (Torch Dynamo _get_handlers() + # asserts on duplicate registration when cuda aliases xpu directly). torch.cuda.Stream = torch.xpu.Stream - torch.cuda.default_stream = torch.xpu.current_stream - torch.cuda.current_stream = torch.xpu.current_stream - torch.cuda.stream = torch.xpu.stream - torch.cuda.set_stream = torch.xpu.set_stream + torch.cuda.default_stream = partial(torch.xpu.current_stream) + torch.cuda.current_stream = partial(torch.xpu.current_stream) + torch.cuda.stream = partial(torch.xpu.stream) + torch.cuda.set_stream = partial(torch.xpu.set_stream) + torch.cuda.Event = partial(torch.xpu.Event) if supports_xpu_graph(): - torch.cuda.graph = torch.xpu.graph + torch.cuda.graph = partial(torch.xpu.graph) torch.cuda.CUDAGraph = torch.xpu.XPUGraph - torch.cuda.graph_pool_handle = torch.xpu.graph_pool_handle + torch.cuda.graph_pool_handle = partial(torch.xpu.graph_pool_handle) yield